paimon 0.3.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Integration tests for REST API.
//!
//! These tests use a mock server to verify the REST API client behavior.
//! Both the mock server and API client run asynchronously using tokio.

use std::collections::HashMap;

use paimon::api::auth::{DLFECSTokenLoader, DLFToken, DLFTokenLoader};
use paimon::api::rest_api::RESTApi;
use paimon::api::ConfigResponse;
use paimon::catalog::{Function, FunctionDefinition, Identifier, ViewSchema};
use paimon::common::Options;
use paimon::spec::DataField;
use serde_json::json;

mod mock_server;
use mock_server::{start_mock_server, RESTServer};
/// Helper struct to hold test resources.
struct TestContext {
    server: RESTServer,
    api: RESTApi,
    url: String,
}

/// Helper function to set up a test environment with a custom prefix.
async fn setup_test_server(initial_dbs: Vec<&str>) -> TestContext {
    let prefix = "mock-test";
    // Create config with prefix
    let mut defaults = HashMap::new();
    defaults.insert("prefix".to_string(), prefix.to_string());
    let config = ConfigResponse::new(defaults);

    let initial: Vec<String> = initial_dbs.iter().map(|s| s.to_string()).collect();
    // Start server with config
    let server = start_mock_server(
        "test_warehouse".to_string(),      // warehouse
        "/tmp/test_warehouse".to_string(), // data_path
        config,
        initial,
    )
    .await;
    let token = "test_token";
    let url = server.url().expect("Failed to get server URL");
    let mut options = Options::new();
    options.set("uri", &url);
    options.set("warehouse", "test_warehouse");
    options.set("token.provider", "bear");
    options.set("token", token);

    let api = RESTApi::new(options, true)
        .await
        .expect("Failed to create RESTApi");

    TestContext { server, api, url }
}

// ==================== Database Tests ====================
#[tokio::test]
async fn test_list_databases() {
    let ctx = setup_test_server(vec!["default", "test_db1", "prod_db"]).await;

    let dbs = ctx.api.list_databases().await.unwrap();

    assert!(dbs.contains(&"default".to_string()));
    assert!(dbs.contains(&"test_db1".to_string()));
    assert!(dbs.contains(&"prod_db".to_string()));
}

#[tokio::test]
async fn test_get_view() {
    let ctx = setup_test_server(vec!["default"]).await;
    let schema: ViewSchema = serde_json::from_value(json!({
        "fields": [{"id": 0, "name": "id", "type": "INT"}],
        "query": "SELECT id FROM source",
        "dialects": {"datafusion": "SELECT id FROM source WHERE id > 0"},
        "comment": null,
        "options": {}
    }))
    .unwrap();
    ctx.server.add_view("default", "active_ids", schema);

    let response = ctx
        .api
        .get_view(&Identifier::new("default", "active_ids"))
        .await
        .unwrap();

    assert_eq!(response.name.as_deref(), Some("active_ids"));
    assert_eq!(response.id.as_deref(), Some("active_ids"));
    assert_eq!(
        response.schema.query_for("datafusion"),
        "SELECT id FROM source WHERE id > 0"
    );
}

#[tokio::test]
async fn test_create_view() {
    let ctx = setup_test_server(vec!["default"]).await;
    let schema = ViewSchema::new(
        serde_json::from_value(json!([
            {"id": 0, "name": "id", "type": "INT"}
        ]))
        .unwrap(),
        "SELECT id FROM source".to_string(),
        HashMap::from([(
            "datafusion".to_string(),
            "SELECT id FROM source".to_string(),
        )]),
        None,
        HashMap::new(),
    );
    let identifier = Identifier::new("default", "active_ids");

    ctx.api.create_view(&identifier, schema).await.unwrap();

    let response = ctx.api.get_view(&identifier).await.unwrap();
    assert_eq!(response.schema.query(), "SELECT id FROM source");
}

#[tokio::test]
async fn test_drop_view_uses_encoded_individual_view_path() {
    let database = "db+%";
    let view = "active+?#%";
    let ctx = setup_test_server(vec![database]).await;
    let schema = ViewSchema::new(
        Vec::new(),
        "SELECT 1".to_string(),
        HashMap::new(),
        None,
        HashMap::new(),
    );
    let identifier = Identifier::new(database, view);
    ctx.api.create_view(&identifier, schema).await.unwrap();

    ctx.api.drop_view(&identifier).await.unwrap();

    assert!(matches!(
        ctx.api.get_view(&identifier).await.unwrap_err(),
        paimon::Error::RestApi {
            source: paimon::api::RestError::NoSuchResource { .. }
        }
    ));
    assert!(matches!(
        ctx.api.drop_view(&identifier).await.unwrap_err(),
        paimon::Error::RestApi {
            source: paimon::api::RestError::NoSuchResource { .. }
        }
    ));
}

#[tokio::test]
async fn test_list_views() {
    let ctx = setup_test_server(vec!["default"]).await;
    ctx.server.set_list_page_size(1);
    let schema: ViewSchema = serde_json::from_value(json!({
        "fields": [{"id": 0, "name": "id", "type": "INT"}],
        "query": "SELECT 1 AS id",
        "dialects": {},
        "comment": null,
        "options": {}
    }))
    .unwrap();
    ctx.server.add_view("default", "v2", schema.clone());
    ctx.server.add_view("default", "v1", schema);

    assert_eq!(
        ctx.api.list_views("default").await.unwrap(),
        vec!["v1", "v2"]
    );
}

#[tokio::test]
async fn test_get_function() {
    let ctx = setup_test_server(vec!["default"]).await;
    let input_params: Vec<DataField> = serde_json::from_value(json!([
        {"id": 0, "name": "length", "type": "DOUBLE"},
        {"id": 1, "name": "width", "type": "DOUBLE"}
    ]))
    .unwrap();
    let return_params: Vec<DataField> = serde_json::from_value(json!([
        {"id": 0, "name": "area", "type": "DOUBLE"}
    ]))
    .unwrap();
    ctx.server.add_function(Function::new(
        Identifier::new("default", "area"),
        Some(input_params),
        Some(return_params),
        true,
        HashMap::from([(
            "datafusion".to_string(),
            FunctionDefinition::Sql {
                definition: "length * width".to_string(),
            },
        )]),
        None,
        HashMap::new(),
    ));

    let response = ctx
        .api
        .get_function(&Identifier::new("default", "area"))
        .await
        .unwrap();

    assert_eq!(response.name.as_deref(), Some("area"));
    assert_eq!(
        response
            .definitions
            .get("datafusion")
            .and_then(FunctionDefinition::sql),
        Some("length * width")
    );
}

#[tokio::test]
async fn test_list_functions() {
    let ctx = setup_test_server(vec!["default"]).await;
    ctx.server.set_list_page_size(1);
    for name in ["zeta", "alpha"] {
        ctx.server.add_function(Function::new(
            Identifier::new("default", name),
            Some(Vec::new()),
            Some(
                serde_json::from_value(json!([
                    {"id": 0, "name": "result", "type": "INT"}
                ]))
                .unwrap(),
            ),
            true,
            HashMap::from([(
                "datafusion".to_string(),
                FunctionDefinition::Sql {
                    definition: "1".to_string(),
                },
            )]),
            None,
            HashMap::new(),
        ));
    }

    assert_eq!(
        ctx.api.list_functions("default").await.unwrap(),
        vec!["alpha", "zeta"]
    );
}

#[tokio::test]
async fn test_create_function() {
    let ctx = setup_test_server(vec!["default"]).await;
    let function = Function::new(
        Identifier::new("default", "plus_one"),
        Some(
            serde_json::from_value(json!([
                {"id": 0, "name": "x", "type": "BIGINT"}
            ]))
            .unwrap(),
        ),
        Some(
            serde_json::from_value(json!([
                {"id": 0, "name": "result", "type": "BIGINT"}
            ]))
            .unwrap(),
        ),
        true,
        HashMap::from([(
            "datafusion".to_string(),
            FunctionDefinition::Sql {
                definition: "x + 1".to_string(),
            },
        )]),
        None,
        HashMap::new(),
    );

    ctx.api.create_function(&function).await.unwrap();

    let response = ctx.api.get_function(function.identifier()).await.unwrap();
    assert_eq!(response.name.as_deref(), Some("plus_one"));
    assert_eq!(response.input_params.unwrap()[0].name(), "x");
}

#[tokio::test]
async fn test_create_database() {
    let ctx = setup_test_server(vec!["default"]).await;

    // Create new database
    let result = ctx.api.create_database("new_db", None).await;
    assert!(result.is_ok(), "failed to create database: {result:?}");

    // Verify creation
    let dbs = ctx.api.list_databases().await.unwrap();
    assert!(dbs.contains(&"new_db".to_string()));

    // Duplicate creation should fail
    let result = ctx.api.create_database("new_db", None).await;
    assert!(result.is_err(), "creating duplicate database should fail");
}

#[tokio::test]
async fn test_get_database() {
    let ctx = setup_test_server(vec!["default"]).await;

    let db_resp = ctx.api.get_database("default").await.unwrap();
    assert_eq!(db_resp.name, Some("default".to_string()));
}

#[tokio::test]
async fn test_error_responses_status_mapping() {
    let ctx = setup_test_server(vec!["default"]).await;

    // Add no-permission database
    ctx.server.add_no_permission_database("secret");

    // GET on no-permission database -> 403
    // Use the prefix from config (v1/mock-test)
    let url = format!("{}/v1/mock-test/databases/{}", ctx.url, "secret");
    let result = reqwest::get(&url).await;
    match result {
        Ok(resp) => {
            assert_eq!(resp.status(), 403);
            let j: serde_json::Value = resp.json().await.unwrap();
            assert_eq!(
                j.get("resourceType").and_then(|v| v.as_str()),
                Some("database")
            );
            assert_eq!(
                j.get("resourceName").and_then(|v| v.as_str()),
                Some("secret")
            );
            assert_eq!(j.get("code").and_then(|v| v.as_u64()), Some(403));
        }
        Err(e) => panic!("Expected 403 response, got error: {e:?}"),
    }

    // POST create existing database -> 409
    let body = json!({"name": "default", "properties": {}});
    let client = reqwest::Client::new();
    let resp = client
        .post(format!("{}/v1/mock-test/databases", ctx.url))
        .json(&body)
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 409);

    let j2: serde_json::Value = resp.json().await.unwrap();
    assert_eq!(
        j2.get("resourceType").and_then(|v| v.as_str()),
        Some("database")
    );
    assert_eq!(
        j2.get("resourceName").and_then(|v| v.as_str()),
        Some("default")
    );
    assert_eq!(j2.get("code").and_then(|v| v.as_u64()), Some(409));
}

#[tokio::test]
async fn test_alter_database() {
    let ctx = setup_test_server(vec!["default"]).await;

    // Alter database with updates
    let mut updates = HashMap::new();
    updates.insert("key1".to_string(), "value1".to_string());
    updates.insert("key2".to_string(), "value2".to_string());

    let result = ctx.api.alter_database("default", vec![], updates).await;
    assert!(result.is_ok(), "failed to alter database: {result:?}");

    // Verify the updates by getting the database
    let db_resp = ctx.api.get_database("default").await.unwrap();
    assert_eq!(db_resp.options.get("key1"), Some(&"value1".to_string()));
    assert_eq!(db_resp.options.get("key2"), Some(&"value2".to_string()));

    // Alter database with removals
    let result = ctx
        .api
        .alter_database("default", vec!["key1".to_string()], HashMap::new())
        .await;
    assert!(result.is_ok(), "failed to remove key: {result:?}");

    let db_resp = ctx.api.get_database("default").await.unwrap();
    assert!(!db_resp.options.contains_key("key1"));
    assert_eq!(db_resp.options.get("key2"), Some(&"value2".to_string()));
}

#[tokio::test]
async fn test_alter_database_not_found() {
    let ctx = setup_test_server(vec!["default"]).await;

    let result = ctx
        .api
        .alter_database("non_existent", vec![], HashMap::new())
        .await;
    assert!(
        result.is_err(),
        "altering non-existent database should fail"
    );
}

#[tokio::test]
async fn test_drop_database() {
    let ctx = setup_test_server(vec!["default", "to_drop"]).await;

    // Verify database exists
    let dbs = ctx.api.list_databases().await.unwrap();
    assert!(dbs.contains(&"to_drop".to_string()));

    // Drop database
    let result = ctx.api.drop_database("to_drop").await;
    assert!(result.is_ok(), "failed to drop database: {result:?}");

    // Verify database is gone
    let dbs = ctx.api.list_databases().await.unwrap();
    assert!(!dbs.contains(&"to_drop".to_string()));

    // Dropping non-existent database should fail
    let result = ctx.api.drop_database("to_drop").await;
    assert!(
        result.is_err(),
        "dropping non-existent database should fail"
    );
}

#[tokio::test]
async fn test_drop_database_no_permission() {
    let ctx = setup_test_server(vec!["default"]).await;
    ctx.server.add_no_permission_database("secret");

    let result = ctx.api.drop_database("secret").await;
    assert!(
        result.is_err(),
        "dropping no-permission database should fail"
    );
}
// ==================== Table Tests ====================

#[tokio::test]
async fn test_list_tables_and_get_table() {
    let ctx = setup_test_server(vec!["default"]).await;

    // Add tables
    ctx.server.add_table("default", "table1");
    ctx.server.add_table("default", "table2");

    // List tables
    let tables = ctx.api.list_tables("default").await.unwrap();
    assert!(tables.contains(&"table1".to_string()));
    assert!(tables.contains(&"table2".to_string()));

    // Get table
    let table_resp = ctx
        .api
        .get_table(&Identifier::new("default", "table1"))
        .await
        .unwrap();
    assert_eq!(table_resp.id.unwrap_or_default(), "table1");
}

#[tokio::test]
async fn test_get_table_not_found() {
    let ctx = setup_test_server(vec!["default"]).await;

    let result = ctx
        .api
        .get_table(&Identifier::new("default", "non_existent_table"))
        .await;
    assert!(result.is_err(), "getting non-existent table should fail");
}

#[tokio::test]
async fn test_list_tables_empty_database() {
    let ctx = setup_test_server(vec!["default"]).await;

    let tables = ctx.api.list_tables("default").await.unwrap();
    assert!(
        tables.is_empty(),
        "expected empty tables list, got: {tables:?}"
    );
}

#[tokio::test]
async fn test_multiple_databases_with_tables() {
    let ctx = setup_test_server(vec!["db1", "db2"]).await;

    // Add tables to different databases
    ctx.server.add_table("db1", "table1_db1");
    ctx.server.add_table("db1", "table2_db1");
    ctx.server.add_table("db2", "table1_db2");

    // Verify db1 tables
    let tables_db1 = ctx.api.list_tables("db1").await.unwrap();
    assert_eq!(tables_db1.len(), 2);
    assert!(tables_db1.contains(&"table1_db1".to_string()));
    assert!(tables_db1.contains(&"table2_db1".to_string()));

    // Verify db2 tables
    let tables_db2 = ctx.api.list_tables("db2").await.unwrap();
    assert_eq!(tables_db2.len(), 1);
    assert!(tables_db2.contains(&"table1_db2".to_string()));
}

#[tokio::test]
async fn test_create_table() {
    let ctx = setup_test_server(vec!["default"]).await;

    // Create a simple schema using builder
    use paimon::spec::{DataType, Schema};
    let schema = Schema::builder()
        .column("id", DataType::BigInt(paimon::spec::BigIntType::new()))
        .column(
            "name",
            DataType::VarChar(paimon::spec::VarCharType::new(255).unwrap()),
        )
        .build()
        .expect("Failed to build schema");

    let result = ctx
        .api
        .create_table(&Identifier::new("default", "new_table"), schema)
        .await;
    assert!(result.is_ok(), "failed to create table: {result:?}");

    // Verify table exists
    let tables = ctx.api.list_tables("default").await.unwrap();
    assert!(tables.contains(&"new_table".to_string()));

    // Get the table
    let table_resp = ctx
        .api
        .get_table(&Identifier::new("default", "new_table"))
        .await
        .unwrap();
    assert_eq!(table_resp.name, Some("new_table".to_string()));
}

#[tokio::test]
async fn test_drop_table() {
    let ctx = setup_test_server(vec!["default"]).await;

    // Add a table
    ctx.server.add_table("default", "table_to_drop");

    // Verify table exists
    let tables = ctx.api.list_tables("default").await.unwrap();
    assert!(tables.contains(&"table_to_drop".to_string()));

    // Drop table
    let result = ctx
        .api
        .drop_table(&Identifier::new("default", "table_to_drop"))
        .await;
    assert!(result.is_ok(), "failed to drop table: {result:?}");

    // Verify table is gone
    let tables = ctx.api.list_tables("default").await.unwrap();
    assert!(!tables.contains(&"table_to_drop".to_string()));

    // Dropping non-existent table should fail
    let result = ctx
        .api
        .drop_table(&Identifier::new("default", "table_to_drop"))
        .await;
    assert!(result.is_err(), "dropping non-existent table should fail");
}

#[tokio::test]
async fn test_drop_table_no_permission() {
    let ctx = setup_test_server(vec!["default"]).await;
    ctx.server
        .add_no_permission_table("default", "secret_table");

    let result = ctx
        .api
        .drop_table(&Identifier::new("default", "secret_table"))
        .await;
    assert!(result.is_err(), "dropping no-permission table should fail");
}

// ==================== Rename Table Tests ====================

#[tokio::test]
async fn test_rename_table() {
    let ctx = setup_test_server(vec!["default"]).await;

    // Add a table
    ctx.server.add_table("default", "old_table");

    // Rename table
    let result = ctx
        .api
        .rename_table(
            &Identifier::new("default", "old_table"),
            &Identifier::new("default", "new_table"),
        )
        .await;
    assert!(result.is_ok(), "failed to rename table: {result:?}");

    // Verify old table is gone
    let tables = ctx.api.list_tables("default").await.unwrap();
    assert!(!tables.contains(&"old_table".to_string()));

    // Verify new table exists
    assert!(tables.contains(&"new_table".to_string()));

    // Get the renamed table
    let table_resp = ctx
        .api
        .get_table(&Identifier::new("default", "new_table"))
        .await
        .unwrap();
    assert_eq!(table_resp.name, Some("new_table".to_string()));
}

// ==================== Token Loader Tests ====================

#[tokio::test]
async fn test_ecs_loader_token() {
    let prefix = "mock-test";
    let mut defaults = HashMap::new();
    defaults.insert("prefix".to_string(), prefix.to_string());
    let config = ConfigResponse::new(defaults);

    let initial: Vec<String> = vec!["default".to_string()];
    let server = start_mock_server(
        "test_warehouse".to_string(),
        "/tmp/test_warehouse".to_string(),
        config,
        initial,
    )
    .await;

    let role_name = "test_role";
    let token_json = json!({
        "AccessKeyId": "AccessKeyId",
        "AccessKeySecret": "AccessKeySecret",
        "SecurityToken": "AQoDYXdzEJr...<remainder of security token>",
        "Expiration": "2023-12-01T12:00:00Z"
    });

    server.set_ecs_metadata(role_name, token_json.clone());

    let ecs_metadata_url = format!("{}/ram/security-credentials/", server.url().unwrap());

    // Test without role name
    let loader = DLFECSTokenLoader::new(&ecs_metadata_url, None);
    let load_token: DLFToken = loader.load_token().await.unwrap();

    assert_eq!(load_token.access_key_id, "AccessKeyId");
    assert_eq!(load_token.access_key_secret, "AccessKeySecret");
    assert_eq!(
        load_token.security_token,
        Some("AQoDYXdzEJr...<remainder of security token>".to_string())
    );
    assert_eq!(
        load_token.expiration,
        Some("2023-12-01T12:00:00Z".to_string())
    );

    // Test with role name
    let loader_with_role = DLFECSTokenLoader::new(&ecs_metadata_url, Some(role_name.to_string()));
    let token: DLFToken = loader_with_role.load_token().await.unwrap();

    assert_eq!(token.access_key_id, "AccessKeyId");
    assert_eq!(token.access_key_secret, "AccessKeySecret");
    assert_eq!(
        token.security_token,
        Some("AQoDYXdzEJr...<remainder of security token>".to_string())
    );
    assert_eq!(token.expiration, Some("2023-12-01T12:00:00Z".to_string()));
}