diff --git a/src/catalog/rest/iceberg_catalog.cpp b/src/catalog/rest/iceberg_catalog.cpp index 79597ee32..0c3600151 100644 --- a/src/catalog/rest/iceberg_catalog.cpp +++ b/src/catalog/rest/iceberg_catalog.cpp @@ -18,6 +18,7 @@ #include "iceberg_logging.hpp" #include "catalog/rest/api/api_utils.hpp" #include "catalog/rest/storage/iceberg_authorization.hpp" +#include "catalog/rest/iceberg_prepare_state.hpp" #include "catalog/rest/storage/authorization/oauth2.hpp" #include "catalog/rest/storage/authorization/sigv4.hpp" #include "catalog/rest/storage/authorization/none.hpp" @@ -519,6 +520,8 @@ void IcebergCatalog::SetAWSCatalogOptions(IcebergAttachOptions &attach_options, unique_ptr IcebergCatalog::Attach(optional_ptr storage_info, ClientContext &context, AttachedDatabase &db, const string &name, AttachInfo &info, AttachOptions &options) { + context.registered_state->GetOrCreate(IcebergPrepareContextState::KEY); + IcebergAttachOptions attach_options; attach_options.warehouse = info.path; attach_options.name = name; diff --git a/src/execution/operator/iceberg_insert.cpp b/src/execution/operator/iceberg_insert.cpp index 9ea90178a..a108a3e47 100644 --- a/src/execution/operator/iceberg_insert.cpp +++ b/src/execution/operator/iceberg_insert.cpp @@ -14,8 +14,11 @@ #include "duckdb/planner/expression/bound_reference_expression.hpp" #include "duckdb/planner/expression/bound_constant_expression.hpp" #include "duckdb/common/multi_file/multi_file_reader.hpp" +#include "duckdb/common/constants.hpp" +#include "duckdb/execution/operator/scan/physical_empty_result.hpp" #include "catalog/rest/iceberg_catalog.hpp" +#include "catalog/rest/iceberg_prepare_state.hpp" #include "catalog/rest/catalog_entry/table/iceberg_table_entry.hpp" #include "execution/operator/iceberg_delete.hpp" #include "catalog/rest/catalog_entry/table/iceberg_table_information.hpp" @@ -30,6 +33,10 @@ namespace duckdb { +static bool IsPrepareOnlyPlanning(ClientContext &context) { + return context.transaction.GetActiveQuery() == MAXIMUM_QUERY_ID; +} + static bool WriteRowId(IcebergInsertVirtualColumns virtual_columns) { return virtual_columns == IcebergInsertVirtualColumns::WRITE_ROW_ID || virtual_columns == IcebergInsertVirtualColumns::WRITE_ROW_ID_AND_SEQUENCE_NUMBER; @@ -934,6 +941,13 @@ PhysicalOperator &IcebergCatalog::PlanCreateTableAs(ClientContext &context, Phys auto &catalog = ic_schema_entry.catalog; auto transaction = catalog.GetCatalogTransaction(context); + if (IsPrepareOnlyPlanning(context)) { + // CTAS has catalog side effects. Prepare caches a no-op plan and execution rebinds with an active query. + auto state = context.registered_state->GetOrCreate(IcebergPrepareContextState::KEY); + state->RequireRebindOnPrepare(); + return planner.Make(op.types, op.estimated_cardinality); + } + // create the table. Takes care of committing to rest catalog and getting the metadata location etc. // setting the schema auto table = ic_schema_entry.CreateTable(transaction, context, *op.info); diff --git a/src/include/catalog/rest/iceberg_prepare_state.hpp b/src/include/catalog/rest/iceberg_prepare_state.hpp new file mode 100644 index 000000000..886bf20ce --- /dev/null +++ b/src/include/catalog/rest/iceberg_prepare_state.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include "duckdb/main/client_context_state.hpp" +#include "duckdb/main/prepared_statement_data.hpp" + +namespace duckdb { + +struct IcebergPrepareContextState : public ClientContextState { + static constexpr const char *KEY = "iceberg_prepare_context"; + + bool CanRequestRebind() override { + return true; + } + + void RequireRebindOnPrepare() { + require_rebind_on_prepare = true; + } + + RebindQueryInfo OnFinalizePrepare(ClientContext &, PreparedStatementData &prepared_statement, + PreparedStatementMode mode) override { + if (mode == PreparedStatementMode::PREPARE_ONLY && require_rebind_on_prepare) { + prepared_statement.properties.always_require_rebind = true; + } + require_rebind_on_prepare = false; + return RebindQueryInfo::DO_NOT_REBIND; + } + +private: + bool require_rebind_on_prepare = false; +}; + +} // namespace duckdb diff --git a/test/python/test_prepared_ctas.py b/test/python/test_prepared_ctas.py new file mode 100644 index 000000000..0db4af1ad --- /dev/null +++ b/test/python/test_prepared_ctas.py @@ -0,0 +1,168 @@ +import ctypes +import glob +import os +import pathlib + +import pytest + + +# SQL PREPARE does not accept CTAS, but ADBC prepares and executes CTAS through +# DuckDB's C API. Exercise that path directly to cover issue #595 without dbt. +pytestmark = pytest.mark.skipif( + os.getenv("ICEBERG_SERVER_AVAILABLE") is None, + reason="Iceberg test catalog is not available", +) + + +def _load_duckdb_library(): + repo = pathlib.Path(__file__).resolve().parents[2] + candidates = [] + for build_type in ("release", "debug"): + candidates.extend(glob.glob(str(repo / "build" / build_type / "src" / "libduckdb.*"))) + candidates = [path for path in candidates if pathlib.Path(path).suffix in (".so", ".dylib", ".dll")] + if not candidates: + pytest.skip("libduckdb was not built") + return ctypes.CDLL(candidates[0]) + + +class DuckDBResult(ctypes.Structure): + _fields_ = [ + ("deprecated_column_count", ctypes.c_uint64), + ("deprecated_row_count", ctypes.c_uint64), + ("deprecated_rows_changed", ctypes.c_uint64), + ("deprecated_columns", ctypes.c_void_p), + ("deprecated_error_message", ctypes.c_char_p), + ("internal_data", ctypes.c_void_p), + ] + + +def _init_api(lib): + lib.duckdb_open.argtypes = [ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)] + lib.duckdb_open.restype = ctypes.c_int + lib.duckdb_connect.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p)] + lib.duckdb_connect.restype = ctypes.c_int + lib.duckdb_disconnect.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + lib.duckdb_close.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + lib.duckdb_query.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(DuckDBResult)] + lib.duckdb_query.restype = ctypes.c_int + lib.duckdb_result_error.argtypes = [ctypes.POINTER(DuckDBResult)] + lib.duckdb_result_error.restype = ctypes.c_char_p + lib.duckdb_destroy_result.argtypes = [ctypes.POINTER(DuckDBResult)] + lib.duckdb_prepare.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)] + lib.duckdb_prepare.restype = ctypes.c_int + lib.duckdb_prepare_error.argtypes = [ctypes.c_void_p] + lib.duckdb_prepare_error.restype = ctypes.c_char_p + lib.duckdb_execute_prepared_streaming.argtypes = [ctypes.c_void_p, ctypes.POINTER(DuckDBResult)] + lib.duckdb_execute_prepared_streaming.restype = ctypes.c_int + lib.duckdb_destroy_prepare.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + lib.duckdb_value_int32.argtypes = [ctypes.POINTER(DuckDBResult), ctypes.c_uint64, ctypes.c_uint64] + lib.duckdb_value_int32.restype = ctypes.c_int32 + lib.duckdb_value_varchar.argtypes = [ctypes.POINTER(DuckDBResult), ctypes.c_uint64, ctypes.c_uint64] + lib.duckdb_value_varchar.restype = ctypes.c_void_p + lib.duckdb_free.argtypes = [ctypes.c_void_p] + + +class DuckDB: + def __init__(self): + self.lib = _load_duckdb_library() + _init_api(self.lib) + self.db = ctypes.c_void_p() + self.con = ctypes.c_void_p() + assert self.lib.duckdb_open(None, ctypes.byref(self.db)) == 0 + assert self.lib.duckdb_connect(self.db, ctypes.byref(self.con)) == 0 + + def close(self): + if self.con: + self.lib.duckdb_disconnect(ctypes.byref(self.con)) + if self.db: + self.lib.duckdb_close(ctypes.byref(self.db)) + + def query(self, sql, expect_ok=True): + result = DuckDBResult() + state = self.lib.duckdb_query(self.con, sql.encode(), ctypes.byref(result)) + error = self._result_error(result) + self.lib.duckdb_destroy_result(ctypes.byref(result)) + if expect_ok: + assert state == 0, error + return state, error + + def prepare(self, sql): + prepared = ctypes.c_void_p() + state = self.lib.duckdb_prepare(self.con, sql.encode(), ctypes.byref(prepared)) + error = self.lib.duckdb_prepare_error(prepared) + assert state == 0, error.decode() if error else None + return prepared + + def execute_prepared_streaming(self, prepared): + result = DuckDBResult() + state = self.lib.duckdb_execute_prepared_streaming(prepared, ctypes.byref(result)) + error = self._result_error(result) + self.lib.duckdb_destroy_result(ctypes.byref(result)) + assert state == 0, error + + def destroy_prepare(self, prepared): + self.lib.duckdb_destroy_prepare(ctypes.byref(prepared)) + + def fetch_single_row(self, sql): + result = DuckDBResult() + state = self.lib.duckdb_query(self.con, sql.encode(), ctypes.byref(result)) + error = self._result_error(result) + assert state == 0, error + text_ptr = self.lib.duckdb_value_varchar(ctypes.byref(result), 1, 0) + assert text_ptr + try: + row = (self.lib.duckdb_value_int32(ctypes.byref(result), 0, 0), ctypes.string_at(text_ptr).decode()) + finally: + self.lib.duckdb_free(text_ptr) + self.lib.duckdb_destroy_result(ctypes.byref(result)) + return row + + def _result_error(self, result): + error = self.lib.duckdb_result_error(ctypes.byref(result)) + return error.decode() if error else None + + +@pytest.fixture() +def duckdb_capi(): + db = DuckDB() + try: + for extension in ("core_functions", "parquet", "avro", "httpfs", "iceberg"): + db.query(f"LOAD {extension}") + db.query( + "CREATE SECRET (TYPE S3,KEY_ID 'admin',SECRET 'password',ENDPOINT '127.0.0.1:9000'," + "URL_STYLE 'path',USE_SSL 0); " + "ATTACH '' AS my_datalake (TYPE ICEBERG,CLIENT_ID 'admin',CLIENT_SECRET 'password'," + "ENDPOINT 'http://127.0.0.1:8181'); " + "CREATE SCHEMA IF NOT EXISTS my_datalake.default;" + ) + yield db + finally: + db.close() + + +def test_iceberg_ctas_prepared_statement_rebinds_at_execute(duckdb_capi): + duckdb_capi.query("DROP TABLE IF EXISTS my_datalake.default.ctas_prepared_rebind_595") + prepared = duckdb_capi.prepare( + "CREATE TABLE my_datalake.default.ctas_prepared_rebind_595 AS SELECT 42 AS id, 'prepared' AS note" + ) + try: + duckdb_capi.execute_prepared_streaming(prepared) + finally: + duckdb_capi.destroy_prepare(prepared) + + assert duckdb_capi.fetch_single_row("SELECT id, note FROM my_datalake.default.ctas_prepared_rebind_595") == ( + 42, + "prepared", + ) + duckdb_capi.query("DROP TABLE my_datalake.default.ctas_prepared_rebind_595") + + +def test_iceberg_duplicate_ctas_in_transaction_still_errors(duckdb_capi): + duckdb_capi.query("DROP TABLE IF EXISTS my_datalake.default.ctas_duplicate_guard_595") + duckdb_capi.query("BEGIN") + duckdb_capi.query("CREATE TABLE my_datalake.default.ctas_duplicate_guard_595 AS SELECT 1 AS id") + state, _ = duckdb_capi.query( + "CREATE TABLE my_datalake.default.ctas_duplicate_guard_595 AS SELECT 2 AS id", expect_ok=False + ) + assert state != 0 + duckdb_capi.query("ROLLBACK")