bff: implement first 100 rows endpoint

[rampage644]

use std::collections::HashMap;
use std::sync::Arc;

use datafusion::prelude::*;
use iceberg_catalog_rest::{RestCatalog, RestCatalogConfig};
use iceberg_datafusion::IcebergCatalogProvider;

#[tokio::main]
pub async fn main() {
    println!("Hello, world?");
}

pub async fn get_catalog() -> IcebergCatalogProvider {
    let config = RestCatalogConfig::builder()
        .uri("https://catalog.embucket.com/catalog".to_string())
        .warehouse("primary_wh".to_string())
        .props(HashMap::default())
        .build();

    let catalog = RestCatalog::new(config);
    IcebergCatalogProvider::try_new(Arc::new(catalog))
        .await
        .unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_first_10_rows_feature() {
        let catalog = get_catalog().await;
        let ctx = SessionContext::new();
        ctx.register_catalog("catalog", Arc::new(catalog));

        let provider = ctx.catalog("catalog").unwrap();
        let schemas = provider.schema_names();
        println!("{schemas:?}");
        assert_eq!(schemas.len(), 6);

        let tables = provider.schema("engineering").unwrap().table_names();
        println!("{tables:?}");
        assert_eq!(tables.len(), 9);

        let table_name = "nyc-taxi-demo";
        let records = ctx
            .sql(&format!(
                "SELECT * FROM catalog.engineering.`{table_name}` LIMIT 10;"
            ))
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();
        assert_eq!(records.len(), 1);
        let record = &records[0];
        assert_eq!(record.num_rows(), 10);

        arrow::util::pretty::print_batches(&records).unwrap();
    }
}

[rampage644]

+----------------------+---------------------+---------------------+--------------+--------------+---------+------------------------+
| dispatching_base_num | pickup_datetime     | dropOff_datetime    | PUlocationID | DOlocationID | SR_Flag | Affiliated_base_number |
+----------------------+---------------------+---------------------+--------------+--------------+---------+------------------------+
| B00013               | 2020-05-01T00:20:21 | 2020-05-01T00:55:09 | 264.0        | 264.0        |         | B00013                 |
| B00014               | 2020-05-01T00:33:29 | 2020-05-01T01:02:58 | 264.0        | 264.0        |         | B00014                 |
| B00021               | 2020-05-01T00:55:23 | 2020-05-01T01:13:59 | 129.0        | 260.0        |         | B00021                 |
| B00021               | 2020-05-01T00:23:10 | 2020-05-01T00:32:04 | 173.0        | 173.0        |         | B00021                 |
| B00037               | 2020-05-01T00:27:29 | 2020-05-01T01:04:29 | 264.0        | 72.0         |         | B00037                 |
| B00037               | 2020-05-01T00:07:36 | 2020-05-01T00:29:28 | 264.0        | 76.0         |         | B00037                 |
| B00037               | 2020-05-01T00:51:58 | 2020-05-01T00:58:50 | 264.0        | 71.0         |         | B00037                 |
| B00112               | 2020-05-01T00:12:52 | 2020-05-01T00:48:53 | 264.0        | 106.0        |         | B00112                 |
| B00221               | 2020-05-01T00:35:34 | 2020-05-01T00:35:35 | 264.0        | 125.0        |         | B00221                 |
| B00221               | 2020-05-01T00:04:38 | 2020-05-01T00:15:43 | 264.0        | 261.0        |         | B00221                 |
+----------------------+---------------------+---------------------+--------------+--------------+---------+------------------------+

[rampage644]

mod tests {
    // use datafusion::catalog_common::CatalogProvider;
    use datafusion::error::DataFusionError;
    use datafusion::physical_plan::display::DisplayableExecutionPlan;
    use datafusion::{config, prelude::*, sql};
    use datafusion_iceberg::catalog::catalog::IcebergCatalog;
    use iceberg_rest_catalog::apis::configuration::Configuration;
    use iceberg_rest_catalog::catalog::RestCatalog;
    use iceberg_rust::catalog::bucket::ObjectStoreBuilder;
    use object_store::aws::{AmazonS3Builder, AwsCredential};
    use std::collections::HashMap;
    use url::Url;

    use super::*;

    #[tokio::test]
    async fn merge_parquet_files() {
        panic!("Not implemented");
        // arrow::util::pretty::print_batches(&res).unwrap();
    }

    pub async fn get_catalog() -> Arc<IcebergCatalog> {
        // let config = RestCatalogConfig::builder()
        //     .uri("https://catalog.embucket.com/catalog".to_string())
        //     .warehouse("fourth".to_string())
        //     .props(HashMap::default())
        //     .build();

        // let client = RestCatalog::new(config);
        // let client = Arc::new(client);
        // Arc::new(IcebergCatalogProvider::try_new(client).await.unwrap())

        let config = {
            let mut config = Configuration::new();
            config.base_path = "https://catalog.embucket.com/catalog".to_string();
            config
        };
        let builder = {
            AmazonS3Builder::from_env()
                .with_bucket_name("embucket-lakehouse")
                .with_region("us-east-2")
                .with_access_key_id("*************")
                .with_secret_access_key("*******************")
        };
        let rest_client = RestCatalog::new(
            Some("9cfed2fc-670f-11ef-b8af-079116b6dad9"),
            config,
            ObjectStoreBuilder::S3(builder),
        );
        let catalog = IcebergCatalog::new(Arc::new(rest_client), None)
            .await
            .unwrap();
        Arc::new(catalog)
    }

    async fn count(ctx: &SessionContext, table: &str) -> usize {
        let res = ctx
            .sql(&format!("SELECT id, salary FROM catalog.sales.`{table}`;"))
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();

        // FIXME: Restore pretty printing
        arrow::util::pretty::print_batches(res.as_slice()).unwrap();

        res[0].num_rows()
    }

    fn get_object_store() -> Arc<dyn object_store::ObjectStore> {
        let region = "us-east-2";
        let bucket_name = "embucket-lakehouse";
        let mut builder = AmazonS3Builder::from_env();

        builder = builder.with_bucket_name(bucket_name);
        builder = builder.with_region(region);
        builder = builder.with_access_key_id("**********");
        builder = builder.with_secret_access_key("************");

        Arc::new(builder.build().unwrap())
    }

    #[tokio::test]
    async fn test_urlencode() {
        let url = "9cfed2fc-670f-11ef-b8af-079116b6dad9";
        assert_eq!(iceberg_rest_catalog::apis::urlencode(url), url);
    }

    #[tokio::test]
    async fn test_insert_values_into_table() {
        let catalog = get_catalog().await;
        let ctx = SessionContext::new();
        let object_store = get_object_store();
        let url = Url::parse("s3://embucket-lakehouse").unwrap();
        ctx.register_catalog("catalog", catalog.clone());
        ctx.register_object_store(&url, object_store);

        let table_name = "sales_sales_transactions_2";

        assert_eq!(count(&ctx, table_name).await, 10);

        let plan = ctx
            .sql(&format!(
                "INSERT INTO catalog.sales.`{table_name}` VALUES ('foo', 'foo', 32, '1987-05-30', 23.0);"
            ))
            .await
            .unwrap();

        let records = plan.collect().await.unwrap();

        assert_eq!(records.len(), 1);

        assert_eq!(count(&ctx, table_name).await, 11);
    }

    #[tokio::test]
    async fn test_copy_csv_to_table() {
        panic!("Not implemented");
    }

    #[tokio::test]
    async fn test_first_10_rows_feature() {
        let catalog = get_catalog().await;
        let ctx = SessionContext::new();
        ctx.register_catalog("catalog", catalog);

        let provider = ctx.catalog("catalog").unwrap();
        let schemas = provider.schema_names();
        assert_eq!(schemas.len(), 7);

        let tables = provider.schema("sales").unwrap().table_names();
        assert_eq!(tables.len(), 2);

        let table_name = "sales_sales_transactions_2";
        let records = ctx
            .sql(&format!(
                "SELECT * FROM catalog.sales.`{table_name}` LIMIT 10;"
            ))
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();
        assert_eq!(records.len(), 1);
        let record = &records[0];
        assert_eq!(record.num_rows(), 10);
    }
}