diff --git a/crates/sail-catalog-iceberg/src/provider.rs b/crates/sail-catalog-iceberg/src/provider.rs index 95397ee06e..2a1069f3be 100644 --- a/crates/sail-catalog-iceberg/src/provider.rs +++ b/crates/sail-catalog-iceberg/src/provider.rs @@ -379,27 +379,36 @@ impl IcebergRestCatalogProvider { catalog_config: &CatalogConfig<'_>, result: &crate::r#gen::LoadTableResult, ) -> CatalogResult<()> { - let rest_session = - Self::rest_table_session_ref(catalog, database, table, catalog_config, result)?; - let mut requirements = Vec::new(); - if rest_session.scan_planning_mode.as_deref() == Some("server") { - requirements.push("server-side scan planning"); + let scan_planning_mode = + Self::effective_scan_planning_mode(result.config.as_ref(), catalog_config)?; + if scan_planning_mode.as_deref() == Some("server") { + return Err(CatalogError::UnsupportedCapability( + "Iceberg REST access session requirements returned by create_table are not supported for create+write yet: server-side scan planning".to_string(), + )); } - if rest_session.remote_signing_enabled { - requirements.push("remote signing"); + + let mut configured_storage_fallbacks = Vec::new(); + if Self::remote_signing_enabled(result.config.as_ref(), catalog_config) { + configured_storage_fallbacks.push("remote signing"); } - if rest_session.storage_credential_count > 0 { - requirements.push("vended credentials"); + if result + .storage_credentials + .as_ref() + .is_some_and(|credentials| !credentials.is_empty()) + { + configured_storage_fallbacks.push("vended credentials"); } - - if requirements.is_empty() { - Ok(()) - } else { - Err(CatalogError::UnsupportedCapability(format!( - "Iceberg REST access session requirements returned by create_table are not supported for create+write yet: {}", - requirements.join(", ") - ))) + if !configured_storage_fallbacks.is_empty() { + log::warn!( + "Iceberg REST catalog {} create_table for {}.{} returned {}; using configured object-store credentials for create+write", + catalog, + quote_namespace_if_needed(database), + quote_name_if_needed(table), + configured_storage_fallbacks.join(", "), + ); } + + Ok(()) } /// Converts an Iceberg REST API table load result into a catalog `TableStatus`. @@ -2039,7 +2048,7 @@ mod tests { } } - fn create_table_response_with_access_session_requirements() -> serde_json::Value { + fn create_table_response_with_access_session_hints() -> serde_json::Value { serde_json::json!({ "metadata-location": "s3://bucket/table/metadata/v1.metadata.json", "metadata": { @@ -2063,7 +2072,6 @@ mod tests { ] }, "config": { - "scan-planning-mode": "server", "s3.remote-signing-enabled": "true" }, "storage-credentials": [ @@ -2078,6 +2086,12 @@ mod tests { }) } + fn create_table_response_with_server_side_scan_planning() -> serde_json::Value { + let mut result = create_table_response_with_access_session_hints(); + result["config"]["scan-planning-mode"] = serde_json::json!("server"); + result + } + async fn load_merged_test_config( defaults: HashMap, mut client_props: HashMap, @@ -2957,13 +2971,33 @@ mod tests { } #[tokio::test] - async fn create_table_rejects_rest_access_session_requirements() { + async fn create_table_allows_rest_access_session_hints() { + let ctx = TestContext::new(Some("test")).await; + let namespace = Namespace::try_from(vec!["db1".to_string()]).unwrap(); + + ctx.mock_post_json( + &ctx.path("/namespaces/db1/tables"), + create_table_response_with_access_session_hints(), + ) + .await; + + let status = ctx + .catalog + .create_table(&namespace, "table1", simple_create_table_options()) + .await + .unwrap(); + + assert_eq!(status.name, "table1"); + } + + #[tokio::test] + async fn create_table_rejects_server_side_scan_planning() { let ctx = TestContext::new(Some("test")).await; let namespace = Namespace::try_from(vec!["db1".to_string()]).unwrap(); ctx.mock_post_json( &ctx.path("/namespaces/db1/tables"), - create_table_response_with_access_session_requirements(), + create_table_response_with_server_side_scan_planning(), ) .await; @@ -2974,17 +3008,17 @@ mod tests { .unwrap_err(); assert!(matches!(err, CatalogError::UnsupportedCapability(_))); - assert!(err.to_string().contains("Iceberg REST access session")); + assert!(err.to_string().contains("server-side scan planning")); } #[tokio::test] - async fn metadata_only_create_table_allows_rest_access_session_requirements() { + async fn metadata_only_create_table_allows_server_side_scan_planning() { let ctx = TestContext::new(Some("test")).await; let namespace = Namespace::try_from(vec!["db1".to_string()]).unwrap(); ctx.mock_post_json( &ctx.path("/namespaces/db1/tables"), - create_table_response_with_access_session_requirements(), + create_table_response_with_server_side_scan_planning(), ) .await; diff --git a/python/pysail/tests/spark/catalog/iceberg_rest/conftest.py b/python/pysail/tests/spark/catalog/iceberg_rest/conftest.py index 5e9da1ec8c..3ad9304122 100644 --- a/python/pysail/tests/spark/catalog/iceberg_rest/conftest.py +++ b/python/pysail/tests/spark/catalog/iceberg_rest/conftest.py @@ -10,8 +10,10 @@ from typing import TYPE_CHECKING import pytest +import requests from testcontainers.core.container import DockerContainer from testcontainers.core.network import Network +from testcontainers.core.wait_strategies import LogMessageWaitStrategy from testcontainers.core.waiting_utils import wait_for_logs from pysail.testing.spark.session import spark_connect_server @@ -20,6 +22,9 @@ from collections.abc import Generator from pathlib import Path +LAKEKEEPER_IMAGE = "quay.io/lakekeeper/catalog:v0.12.1" +LAKEKEEPER_DATABASE_URL = "postgresql://postgres:postgres@lakekeeper-db:5432/postgres" +LAKEKEEPER_PROJECT_ID = "00000000-0000-0000-0000-000000000000" NESSIE_NAMESPACE_SEPARATOR = "-" @@ -54,9 +59,9 @@ def seaweedfs_container( .with_exposed_ports(8333) .with_network(docker_network) .with_network_aliases("seaweedfs") + .waiting_for(LogMessageWaitStrategy("Start Seaweed S3 API").with_startup_timeout(120)) ) container.start() - wait_for_logs(container, "Start Seaweed S3 API", timeout=120) yield container container.stop() @@ -102,6 +107,143 @@ def _create_s3_bucket(seaweedfs_host_endpoint: str) -> None: return +def lakekeeper_command_container( + docker_network: Network, + command: str, +) -> DockerContainer: + """Configure a Lakekeeper container for the test PostgreSQL database.""" + return ( + DockerContainer(LAKEKEEPER_IMAGE) + .with_command([command]) + .with_env("LAKEKEEPER__PG_ENCRYPTION_KEY", "This-is-NOT-Secure!") + .with_env("LAKEKEEPER__PG_DATABASE_URL_READ", LAKEKEEPER_DATABASE_URL) + .with_env("LAKEKEEPER__PG_DATABASE_URL_WRITE", LAKEKEEPER_DATABASE_URL) + .with_network(docker_network) + ) + + +@pytest.fixture(scope="module") +def lakekeeper_database_container( + docker_network: Network, +) -> Generator[DockerContainer, None, None]: + """Start the PostgreSQL database used by Lakekeeper.""" + container = ( + DockerContainer("postgres:17") + .with_env("POSTGRES_PASSWORD", "postgres") + .with_network(docker_network) + .with_network_aliases("lakekeeper-db") + .waiting_for(LogMessageWaitStrategy("database system is ready to accept connections").with_startup_timeout(120)) + ) + container.start() + try: + yield container + finally: + container.stop() + + +@pytest.fixture(scope="module") +def _lakekeeper_migration( + docker_network: Network, + lakekeeper_database_container: DockerContainer, # noqa: ARG001 +) -> None: + """Apply Lakekeeper database migrations before starting the server.""" + container = lakekeeper_command_container(docker_network, "migrate") + container.start() + try: + result = container.get_wrapped_container().wait(timeout=120) + if result["StatusCode"] != 0: + stdout, stderr = container.get_logs() + message = (stdout + stderr).decode(errors="replace") + msg = f"Lakekeeper database migration failed:\n{message}" + raise RuntimeError(msg) + finally: + container.stop() + + +@pytest.fixture(scope="module") +def lakekeeper_container( + docker_network: Network, + seaweedfs_container: DockerContainer, # noqa: ARG001 + _lakekeeper_migration: None, +) -> Generator[DockerContainer, None, None]: + """Start Lakekeeper after its database has been migrated.""" + container = ( + lakekeeper_command_container(docker_network, "serve") + .with_exposed_ports(8181) + .with_network_aliases("lakekeeper") + .waiting_for(LogMessageWaitStrategy("Starting server on 0.0.0.0:8181").with_startup_timeout(120)) + ) + container.start() + try: + yield container + finally: + container.stop() + + +@pytest.fixture(scope="module") +def lakekeeper_endpoint(lakekeeper_container: DockerContainer) -> str: + """Return a healthy host-accessible Lakekeeper endpoint.""" + host = lakekeeper_container.get_container_host_ip() + port = lakekeeper_container.get_exposed_port(8181) + endpoint = f"http://{host}:{port}" + + for attempt in range(30): + try: + response = requests.get(f"{endpoint}/health", timeout=10) + response.raise_for_status() + except requests.RequestException: + if attempt == 29: # noqa: PLR2004 + raise + time.sleep(1) + else: + return endpoint + msg = "unreachable" + raise AssertionError(msg) + + +@pytest.fixture(scope="module") +def lakekeeper_warehouse_id( + lakekeeper_endpoint: str, + seaweedfs_internal_endpoint: str, + _create_s3_bucket: None, +) -> str: + """Bootstrap Lakekeeper and create the S3-compatible test warehouse.""" + bootstrap = requests.post( + f"{lakekeeper_endpoint}/management/v1/bootstrap", + json={"accept-terms-of-use": True}, + timeout=30, + ) + bootstrap.raise_for_status() + + warehouse = requests.post( + f"{lakekeeper_endpoint}/management/v1/warehouse", + json={ + "warehouse-name": "demo", + "project-id": LAKEKEEPER_PROJECT_ID, + "storage-profile": { + "type": "s3", + "bucket": "icebergdata", + "key-prefix": "lakekeeper", + "endpoint": seaweedfs_internal_endpoint, + "region": "us-east-1", + "path-style-access": True, + "flavor": "s3-compat", + "sts-enabled": False, + "remote-signing-enabled": True, + }, + "storage-credential": { + "type": "s3", + "credential-type": "access-key", + "access-key-id": "admin", + "secret-access-key": "password", + }, + }, + timeout=30, + ) + warehouse.raise_for_status() + return warehouse.json()["warehouse-id"] + + @pytest.fixture(scope="module") def iceberg_rest_container( docker_network: Network, diff --git a/python/pysail/tests/spark/catalog/iceberg_rest/test_lakekeeper.py b/python/pysail/tests/spark/catalog/iceberg_rest/test_lakekeeper.py new file mode 100644 index 0000000000..c8e1073150 --- /dev/null +++ b/python/pysail/tests/spark/catalog/iceberg_rest/test_lakekeeper.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import urllib.parse +from typing import TYPE_CHECKING + +import pytest +import requests + +from pysail.testing.spark.session import spark_connect_server + +if TYPE_CHECKING: + from collections.abc import Generator + + from pyspark.sql import SparkSession + +NAMESPACE = "lakekeeper_access_session_test" +TABLE = f"sail.{NAMESPACE}.remote_signing_t" + + +@pytest.fixture(scope="module") +def remote( + lakekeeper_endpoint: str, + lakekeeper_warehouse_id: str, + seaweedfs_host_endpoint: str, +) -> Generator[str, None, None]: + """Start Sail with the Lakekeeper-backed Iceberg REST catalog.""" + del lakekeeper_warehouse_id + catalog_config = f'[{{name="sail", type="iceberg-rest", uri="{lakekeeper_endpoint}/catalog", warehouse="demo"}}]' + with spark_connect_server( + envs={ + "SAIL_CATALOG__LIST": catalog_config, + "AWS_ACCESS_KEY_ID": "admin", + "AWS_SECRET_ACCESS_KEY": "password", + "AWS_REGION": "us-east-1", + "AWS_ENDPOINT": seaweedfs_host_endpoint, + "AWS_VIRTUAL_HOSTED_STYLE_REQUEST": "false", + "AWS_ALLOW_HTTP": "true", + }, + ) as server: + yield server.remote + + +@pytest.fixture(scope="module", autouse=True) +def namespace(spark: SparkSession) -> Generator[None, None, None]: + spark.sql(f"CREATE NAMESPACE IF NOT EXISTS sail.{NAMESPACE}") + yield + spark.sql(f"DROP NAMESPACE IF EXISTS sail.{NAMESPACE} CASCADE") + + +def _load_lakekeeper_table( + lakekeeper_endpoint: str, + lakekeeper_warehouse_id: str, +) -> dict: + namespace = urllib.parse.quote(NAMESPACE, safe="") + table = urllib.parse.quote(TABLE.rsplit(".", 1)[1], safe="") + response = requests.get( + f"{lakekeeper_endpoint}/catalog/v1/{lakekeeper_warehouse_id}/namespaces/{namespace}/tables/{table}", + timeout=30, + ) + response.raise_for_status() + return response.json() + + +def test_create_write_uses_configured_credentials_with_lakekeeper_session_hints( + spark: SparkSession, + lakekeeper_endpoint: str, + lakekeeper_warehouse_id: str, +) -> None: + spark.sql(f"DROP TABLE IF EXISTS {TABLE}") + source = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "name"]) + source.writeTo(TABLE).using("iceberg").create() + + rows = spark.table(TABLE).orderBy("id").collect() + assert [(row["id"], row["name"]) for row in rows] == [(1, "a"), (2, "b")] + + table = _load_lakekeeper_table(lakekeeper_endpoint, lakekeeper_warehouse_id) + assert table["config"]["s3.remote-signing-enabled"] == "true" + assert table["storage-credentials"]