From 296fdb90f012717e3f190a27e83f6fff2a90715d Mon Sep 17 00:00:00 2001 From: Kanthi Subramanian Date: Sun, 18 Oct 2026 23:25:47 +0200 Subject: [PATCH] Merge pull request #115879 from ClickHouse/asya/unity_v2 Add new Unity catalog implementation with Delta Lake and Iceberg support --- src/Databases/DataLake/DatabaseDataLake.cpp | 97 +- src/Databases/DataLake/DatabaseDataLake.h | 3 +- .../DataLake/DatabaseDataLakeSettings.cpp | 15 +- .../DataLake/DatabaseDataLakeSettings.h | 1 + src/Databases/DataLake/GlueCatalog.cpp | 2 +- src/Databases/DataLake/GlueCatalog.h | 2 +- .../DataLake/HTTPBasedCatalogUtils.cpp | 53 + .../DataLake/HTTPBasedCatalogUtils.h | 45 + src/Databases/DataLake/ICatalog.cpp | 19 + src/Databases/DataLake/ICatalog.h | 31 +- src/Databases/DataLake/RestCatalog.cpp | 86 +- src/Databases/DataLake/RestCatalog.h | 17 +- src/Databases/DataLake/UnityCatalog.cpp | 2 +- src/Databases/DataLake/UnityCatalog.h | 2 +- src/Databases/DataLake/UnityCatalogUtils.cpp | 121 +++ src/Databases/DataLake/UnityCatalogUtils.h | 26 + src/Databases/DataLake/UnityV2Catalog.cpp | 915 ++++++++++++++++++ src/Databases/DataLake/UnityV2Catalog.h | 176 ++++ tests/integration/test_database_delta/test.py | 66 +- .../test_database_unity_v2/__init__.py | 1 + .../configs/display_secrets.xml | 3 + .../configs/user_files_root.xml | 7 + .../mock_servers/uc_proxy.py | 199 ++++ .../test_database_unity_v2/test.py | 651 +++++++++++++ 24 files changed, 2407 insertions(+), 133 deletions(-) create mode 100644 src/Databases/DataLake/UnityCatalogUtils.cpp create mode 100644 src/Databases/DataLake/UnityCatalogUtils.h create mode 100644 src/Databases/DataLake/UnityV2Catalog.cpp create mode 100644 src/Databases/DataLake/UnityV2Catalog.h create mode 100644 tests/integration/test_database_unity_v2/__init__.py create mode 100644 tests/integration/test_database_unity_v2/configs/display_secrets.xml create mode 100644 tests/integration/test_database_unity_v2/configs/user_files_root.xml create mode 100644 tests/integration/test_database_unity_v2/mock_servers/uc_proxy.py create mode 100644 tests/integration/test_database_unity_v2/test.py diff --git a/src/Databases/DataLake/DatabaseDataLake.cpp b/src/Databases/DataLake/DatabaseDataLake.cpp index c5c40e3bd9aa5..944ddc675caa7 100644 --- a/src/Databases/DataLake/DatabaseDataLake.cpp +++ b/src/Databases/DataLake/DatabaseDataLake.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include #if USE_AWS_S3 && USE_SSL @@ -71,6 +72,7 @@ namespace DatabaseDataLakeSetting extern const DatabaseDataLakeSettingsDatabaseDataLakeCatalogType catalog_type; extern const DatabaseDataLakeSettingsString warehouse; extern const DatabaseDataLakeSettingsString catalog_credential; + extern const DatabaseDataLakeSettingsBool use_unity_catalog_v2; extern const DatabaseDataLakeSettingsString auth_header; extern const DatabaseDataLakeSettingsString auth_scope; extern const DatabaseDataLakeSettingsString storage_endpoint; @@ -385,11 +387,28 @@ void DatabaseDataLake::initialize() const } case DB::DatabaseDataLakeCatalogType::UNITY: { - catalog_impl = std::make_shared( + if (!settings[DatabaseDataLakeSetting::use_unity_catalog_v2]) + { + catalog_impl = std::make_shared( + settings[DatabaseDataLakeSetting::warehouse].value, + url, + settings[DatabaseDataLakeSetting::catalog_credential].value, + settings[DatabaseDataLakeSetting::namespaces].value, + Context::getGlobalContextInstance()); + break; + } + + /// Databricks OIDC expects `all-apis`; the default `auth_scope` value targets Iceberg REST catalogs. + const std::string unity_auth_scope = settings[DatabaseDataLakeSetting::auth_scope].changed + ? settings[DatabaseDataLakeSetting::auth_scope].value + : "all-apis"; + catalog_impl = std::make_shared( settings[DatabaseDataLakeSetting::warehouse].value, url, settings[DatabaseDataLakeSetting::catalog_credential].value, - settings[DatabaseDataLakeSetting::namespaces].value, + unity_auth_scope, + settings[DatabaseDataLakeSetting::auth_header].value, + settings[DatabaseDataLakeSetting::oauth_server_uri].value, Context::getGlobalContextInstance()); break; } @@ -537,7 +556,8 @@ void DatabaseDataLake::resetCatalog(String reason) const StorageObjectStorageConfigurationPtr DatabaseDataLake::getConfiguration( DatabaseDataLakeStorageType type, - DataLakeStorageSettingsPtr storage_settings) const + DataLakeStorageSettingsPtr storage_settings, + DataLake::DataLakeTableFormat table_format) const { /// TODO: add tests for azure, local storage types. @@ -545,7 +565,14 @@ StorageObjectStorageConfigurationPtr DatabaseDataLake::getConfiguration( const DatabaseDataLakeSettings & settings = *settings_version; auto catalog = getCatalog(); - switch (catalog->getCatalogType()) + auto catalog_type = catalog->getCatalogType(); + + /// `UnityV2Catalog` serves both formats. Its Iceberg tables need the same configuration as an + /// Iceberg REST catalog, so route them to that arm; the legacy catalog never reports Iceberg. + if (catalog_type == DatabaseDataLakeCatalogType::UNITY && table_format == DataLake::DataLakeTableFormat::ICEBERG) + catalog_type = DatabaseDataLakeCatalogType::ICEBERG_REST; + + switch (catalog_type) { case DatabaseDataLakeCatalogType::ICEBERG_ONELAKE: { @@ -890,7 +917,7 @@ StoragePtr DatabaseDataLake::tryGetTableImpl(const String & name, ContextPtr con } } - const auto configuration = getConfiguration(storage_type, storage_settings); + const auto configuration = getConfiguration(storage_type, storage_settings, table_metadata.getTableFormat()); /// HACK: Hacky-hack to enable lazy load ContextMutablePtr context_copy = Context::createCopy(context_); @@ -972,7 +999,7 @@ StoragePtr DatabaseDataLake::tryGetTableImpl(const String & name, ContextPtr con return std::nullopt; if (!with_vended_credentials && !catalog_manages_provider_chain) return std::nullopt; - return catalog->getCredentialsConfigurationCallback(storage_id); + return catalog->getCredentialsConfigurationCallback(storage_id, table_metadata); }; const auto catalog_uuid = table_metadata.getTableUUID(); @@ -1298,7 +1325,7 @@ DatabaseTablesIteratorPtr DatabaseDataLake::getTablesIteratorWithHint( DatabaseTablesIteratorPtr DatabaseDataLake::getTablesIteratorImpl( ContextPtr context_, const FilterByNameFunction & filter_by_table_name, - bool skip_not_loaded, + bool /*skip_not_loaded*/, const TablesFilter & tables_filter, bool keep_unresolved_tables) const { @@ -1347,13 +1374,15 @@ DatabaseTablesIteratorPtr DatabaseDataLake::getTablesIteratorImpl( futures.emplace_back(promises.back()->get_future()); pool.scheduleOrThrow( - [this, table_name, skip_not_loaded, context_, keep_unresolved_tables, promise=promises.back()]() mutable + [this, table_name, context_, keep_unresolved_tables, promise=promises.back()]() mutable { StoragePtr storage = nullptr; try { LOG_INFO(log, "Get table information for table {}", table_name); - storage = tryGetTableImpl(table_name, context_, false, skip_not_loaded); + /// Tables marked unreadable at metadata time are skipped, not an error; + /// direct access still reports the reason. + storage = tryGetTableImpl(table_name, context_, /*lightweight*/ false, /*ignore_if_not_iceberg*/ true); } catch (...) { @@ -1545,6 +1574,10 @@ void DatabaseDataLake::applySettingsChanges(const SettingsChanges & settings_cha auto new_settings = std::make_unique(*current_settings); new_settings->applyChanges(settings_changes); + /// Switching the Unity implementation replaces the catalog object rather than altering it in place. + const bool implementation_changed = (*new_settings)[DatabaseDataLakeSetting::use_unity_catalog_v2].value + != (*current_settings)[DatabaseDataLakeSetting::use_unity_catalog_v2].value; + ASTPtr new_engine_definition; { std::lock_guard lock(mutex); @@ -1572,7 +1605,17 @@ void DatabaseDataLake::applySettingsChanges(const SettingsChanges & settings_cha storage->set(storage->settings, storage_settings_ast); } + /// `use_unity_catalog_v2` is consumed by the database, not by the catalog object, so it is + /// never handed to the catalog. Re-applying the current value is a metadata-only no-op. + SettingsChanges catalog_settings_changes; + for (const auto & change : settings_changes) + { + if (change.name != DatabaseDataLakeSettings::getSettingName(DatabaseDataLakeSetting::use_unity_catalog_v2)) + catalog_settings_changes.push_back(change); + } + std::shared_ptr local_catalog_snapshot; + if (!implementation_changed) { std::lock_guard lock(catalog_mutex); local_catalog_snapshot = catalog_impl; @@ -1581,8 +1624,9 @@ void DatabaseDataLake::applySettingsChanges(const SettingsChanges & settings_cha /// Prepare the new catalog state without publishing it: validation, the eager token /// fetch and the config reload may throw, and then nothing has changed yet. DataLake::ICatalog::PreparedSettingsChangesPtr prepared_catalog_changes; - if (local_catalog_snapshot) - prepared_catalog_changes = local_catalog_snapshot->prepareSettingsChanges(settings_changes); + const bool alter_catalog = local_catalog_snapshot && !catalog_settings_changes.empty(); + if (alter_catalog) + prepared_catalog_changes = local_catalog_snapshot->prepareSettingsChanges(catalog_settings_changes); /// Persist the new metadata before publishing anything: if the write fails, the live /// state is untouched and matches the old metadata on disk. The create query is built @@ -1594,7 +1638,7 @@ void DatabaseDataLake::applySettingsChanges(const SettingsChanges & settings_cha DatabaseCatalog::instance().updateMetadataFile(getDatabaseName(), new_create_query); /// Publish. Nothing below throws. - if (local_catalog_snapshot) + if (alter_catalog) local_catalog_snapshot->commitSettingsChanges(std::move(prepared_catalog_changes)); database_settings.set(std::move(new_settings)); { @@ -1603,9 +1647,9 @@ void DatabaseDataLake::applySettingsChanges(const SettingsChanges & settings_cha } if (!local_catalog_snapshot) { - /// The catalog was not built when the ALTER started. If a concurrent query - /// built it meanwhile, it used the old settings: drop it so the next access - /// rebuilds it with the new ones. Also clear a recorded construction failure + /// The catalog was not built when the ALTER started, or the implementation was switched. + /// If a concurrent query built it meanwhile, it used the old settings: drop it so the next + /// access rebuilds it with the new ones. Also clear a recorded construction failure /// (e.g. credentials lost on RESTORE) for the same reason. std::lock_guard lock(catalog_mutex); resetCatalog(/* reason */ ""); @@ -1690,6 +1734,10 @@ ASTPtr DatabaseDataLake::getCreateTableQueryImpl( storage->engine->setKind(ASTFunction::Kind::TABLE_ENGINE); if (!table_metadata.isDefaultReadableTable()) storage->engine->name = DataLake::FAKE_TABLE_ENGINE_NAME_FOR_UNREADABLE_TABLES; + else if (table_metadata.getTableFormat() != DataLake::DataLakeTableFormat::UNKNOWN) + /// Mixed-format catalogs (e.g. `UnityV2Catalog`) set the format per table; single-format + /// catalogs leave it UNKNOWN and keep the database-level table engine name. + storage->engine->name = String(catalog->getTableEngineName(table_metadata)); storage->settings = {}; if (partition_by) @@ -2046,9 +2094,10 @@ The following settings are supported: | Setting | Description | |-------------------------|-----------------------------------------------------------------------------------------| -| `catalog_type` | Type of catalog: `glue`, `unity` (Delta), `rest` (Iceberg), `hive`, `onelake` (Iceberg), `delta_sharing` (Iceberg, flat namespaces), `horizon` (Snowflake Horizon Iceberg REST) | +| `catalog_type` | Type of catalog: `glue`, `unity` (Delta, or Delta and Iceberg with `use_unity_catalog_v2`), `rest` (Iceberg), `hive`, `onelake` (Iceberg), `delta_sharing` (Iceberg, flat namespaces), `horizon` (Snowflake Horizon Iceberg REST) | | `warehouse` | The warehouse/database name to use in the catalog. | | `catalog_credential` | Authentication credential for the catalog (e.g., API key or token) | +| `use_unity_catalog_v2` | For `catalog_type = 'unity'`: use the new implementation, which serves both Delta Lake and Iceberg tables. Default: `false`. Set it on `CREATE`, or change it for an existing database with `ALTER DATABASE ... MODIFY SETTING`. | | `auth_header` | Custom HTTP header for authentication with the catalog service | | `auth_scope` | OAuth2 scope for authentication (if using OAuth) | | `storage_endpoint` | Endpoint URL for the underlying storage | @@ -2174,6 +2223,22 @@ Some catalogs support only one of the two modes: See below sections for examples of using the `DataLakeCatalog` engine: * [Unity Catalog](/guides/use-cases/data-warehousing/unity-catalog) +* Unity Catalog with Delta Lake and Iceberg tables + With `use_unity_catalog_v2 = 1`, the `unity` catalog serves both Delta Lake and Iceberg tables, + detecting the format of each table. +```sql +CREATE DATABASE database_name +ENGINE = DataLakeCatalog('https://.cloud.databricks.com/api/2.1/unity-catalog') +SETTINGS + catalog_type = 'unity', + use_unity_catalog_v2 = 1, + warehouse = 'my_catalog', + catalog_credential = ':'; +SHOW TABLES FROM database_name; +SELECT count() FROM database_name.`schema_name.table_name`; +``` + `catalog_credential` accepts an OAuth service principal as `:`, + a personal access token (`dapi...`), or a pre-obtained OAuth token. * [Glue Catalog](/guides/use-cases/data-warehousing/glue-catalog) * OneLake Catalog Can be used by enabling `allow_experimental_database_iceberg` or `allow_database_iceberg`. diff --git a/src/Databases/DataLake/DatabaseDataLake.h b/src/Databases/DataLake/DatabaseDataLake.h index e1101c2232146..b98999ca9998a 100644 --- a/src/Databases/DataLake/DatabaseDataLake.h +++ b/src/Databases/DataLake/DatabaseDataLake.h @@ -142,7 +142,8 @@ class DatabaseDataLake final : public IDatabase, WithContext StorageObjectStorageConfigurationPtr getConfiguration( DatabaseDataLakeStorageType type, - DataLakeStorageSettingsPtr storage_settings) const; + DataLakeStorageSettingsPtr storage_settings, + DataLake::DataLakeTableFormat table_format) const; std::string getStorageEndpointForTable(const DataLake::TableMetadata & table_metadata) const; diff --git a/src/Databases/DataLake/DatabaseDataLakeSettings.cpp b/src/Databases/DataLake/DatabaseDataLakeSettings.cpp index 5d20544965221..cde9b7c2977b1 100644 --- a/src/Databases/DataLake/DatabaseDataLakeSettings.cpp +++ b/src/Databases/DataLake/DatabaseDataLakeSettings.cpp @@ -20,6 +20,7 @@ namespace ErrorCodes #define DATABASE_ICEBERG_RELATED_SETTINGS(DECLARE, ALIAS) \ DECLARE(DatabaseDataLakeCatalogType, catalog_type, DatabaseDataLakeCatalogType::NONE, "Catalog type", 0) \ DECLARE(String, catalog_credential, "", "", 0) \ + DECLARE(Bool, use_unity_catalog_v2, false, "Use the new Unity catalog implementation, which serves both Delta Lake and Iceberg tables. Only for catalog_type = 'unity'", 0) \ DECLARE(Bool, vended_credentials, true, "Use vended credentials (storage credentials) from catalog", 0) \ DECLARE(UInt64, vended_credentials_cache_ttl, 300, "Maximum cache entry lifetime (in seconds) for vended credentials. '0' disables caching.", 0) \ DECLARE(String, auth_scope, "PRINCIPAL_ROLE:ALL", "Authorization scope for client credentials or token exchange", 0) \ @@ -113,13 +114,23 @@ SettingsChanges DatabaseDataLakeSettings::allChanged() const return changes; } -const String & DatabaseDataLakeSettings::getSettingName(DatabaseDataLakeSettingsString setting) +static const String & getSettingNameByOffset(size_t offset) { const auto & accessor = DatabaseDataLakeSettingsTraits::Accessor::instance(); - const size_t index = accessor.findByOffset(setting.offset); + const size_t index = accessor.findByOffset(offset); if (index == static_cast(-1)) throw Exception(ErrorCodes::LOGICAL_ERROR, "Unknown database DataLake setting"); return accessor.getName(index); } +const String & DatabaseDataLakeSettings::getSettingName(DatabaseDataLakeSettingsString setting) +{ + return getSettingNameByOffset(setting.offset); +} + +const String & DatabaseDataLakeSettings::getSettingName(DatabaseDataLakeSettingsBool setting) +{ + return getSettingNameByOffset(setting.offset); +} + } diff --git a/src/Databases/DataLake/DatabaseDataLakeSettings.h b/src/Databases/DataLake/DatabaseDataLakeSettings.h index ad94cee911654..a30a580464702 100644 --- a/src/Databases/DataLake/DatabaseDataLakeSettings.h +++ b/src/Databases/DataLake/DatabaseDataLakeSettings.h @@ -75,6 +75,7 @@ struct DatabaseDataLakeSettings /// Name of the setting referenced by its `DatabaseDataLakeSetting::*` index, /// so catalog code can match `SettingsChanges` entries without magic strings. static const String & getSettingName(DatabaseDataLakeSettingsString setting); + static const String & getSettingName(DatabaseDataLakeSettingsBool setting); private: std::unique_ptr impl; diff --git a/src/Databases/DataLake/GlueCatalog.cpp b/src/Databases/DataLake/GlueCatalog.cpp index 764684e19c12d..792e8997b0f69 100644 --- a/src/Databases/DataLake/GlueCatalog.cpp +++ b/src/Databases/DataLake/GlueCatalog.cpp @@ -648,7 +648,7 @@ void GlueCatalog::setCredentials(TableMetadata & metadata) const } } -ICatalog::CredentialsRefreshCallback GlueCatalog::getCredentialsConfigurationCallback(const DB::StorageID & storage_id) +ICatalog::CredentialsRefreshCallback GlueCatalog::getCredentialsConfigurationCallback(const DB::StorageID & storage_id, const TableMetadata & /*table_metadata*/) { /// The AWS SDK credentials provider chain (instance profile, STS assume-role, /// web-identity, etc.) refreshes its cached credentials internally before diff --git a/src/Databases/DataLake/GlueCatalog.h b/src/Databases/DataLake/GlueCatalog.h index 41c7deba8a003..654c7cdf42321 100644 --- a/src/Databases/DataLake/GlueCatalog.h +++ b/src/Databases/DataLake/GlueCatalog.h @@ -97,7 +97,7 @@ class GlueCatalog final : public ICatalog, private DB::WithContext /// credentials provider chain. Invoked by `ReadBufferFromS3` when an S3 call /// fails with `ExpiredToken`, so that a long-running read can recover without /// the user having to restart the query. - ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & storage_id) override; + ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & storage_id, const TableMetadata & table_metadata) override; /// Resolves the precise Iceberg timestamp type for `column_name` by searching the current schema /// in the Iceberg `metadata_object`. Falls back to `"timestamp_ns"` when `glue_column_type` is diff --git a/src/Databases/DataLake/HTTPBasedCatalogUtils.cpp b/src/Databases/DataLake/HTTPBasedCatalogUtils.cpp index 93d4d6313ba12..8f0d926e7f2d0 100644 --- a/src/Databases/DataLake/HTTPBasedCatalogUtils.cpp +++ b/src/Databases/DataLake/HTTPBasedCatalogUtils.cpp @@ -5,6 +5,10 @@ #include #include #include +#include +#include +#include +#include namespace DB::ErrorCodes { @@ -101,5 +105,54 @@ std::pair makeHTTPRequestAndReadJSON( } } +AccessToken requestOAuthToken(const DB::ContextPtr & context, const Poco::URI & url, const std::string & body) +{ + context->getRemoteHostFilter().checkHostAndPort(url.getHost(), std::to_string(url.getPort())); + auto timeouts = DB::ConnectionTimeouts::getHTTPTimeouts(context->getSettingsRef(), context->getServerSettings()); + auto session = makeHTTPSession(DB::HTTPConnectionGroupType::HTTP, url, timeouts, {}); + + Poco::Net::HTTPRequest request(Poco::Net::HTTPRequest::HTTP_POST, url.getPathAndQuery(), Poco::Net::HTTPMessage::HTTP_1_1); + request.setContentType("application/x-www-form-urlencoded"); + request.setContentLength(body.size()); + request.set("Accept", "application/json"); + + session->sendRequest(request) << body; + + Poco::Net::HTTPResponse response; + std::istream & rs = session->receiveResponse(response); + + std::string json_str; + Poco::StreamCopier::copyToString(rs, json_str); + + /// The body of a failed response is an OAuth error object, safe to show. + if (response.getStatus() != Poco::Net::HTTPResponse::HTTP_OK) + throw DB::Exception( + DB::ErrorCodes::DATALAKE_DATABASE_ERROR, + "OAuth token request failed with status {} ({}): {}", + static_cast(response.getStatus()), response.getReason(), json_str); + + Poco::JSON::Parser parser; + Poco::Dynamic::Var res_json = parser.parse(json_str); + const Poco::JSON::Object::Ptr & object = res_json.extract(); + + if (!object->has("access_token")) + throw DB::Exception( + DB::ErrorCodes::DATALAKE_DATABASE_ERROR, + "OAuth token response has no `access_token` field: {}", + json_str); + + AccessToken token; + token.token = object->getValue("access_token"); + + if (object->has("expires_in")) + { + Int64 expires_in = object->getValue("expires_in"); + /// Use 90% of the token lifetime as the validity window so that short-lived tokens + /// (e.g. expires_in=300) still get a sensible buffer instead of going non-positive. + token.expires_at = std::chrono::system_clock::now() + std::chrono::seconds(expires_in * 9 / 10); + } + + return token; +} } diff --git a/src/Databases/DataLake/HTTPBasedCatalogUtils.h b/src/Databases/DataLake/HTTPBasedCatalogUtils.h index 7555a5da7d25f..c571d5d12682d 100644 --- a/src/Databases/DataLake/HTTPBasedCatalogUtils.h +++ b/src/Databases/DataLake/HTTPBasedCatalogUtils.h @@ -1,12 +1,57 @@ #pragma once +#include #include #include +#include +#include #include +#include namespace DataLake { +struct AccessToken +{ + std::string token; + std::optional expires_at; + + bool isExpired() const + { + if (!expires_at.has_value()) + return false; + return std::chrono::system_clock::now() >= expires_at.value(); + } +}; + +/// POSTs an OAuth token request. An empty `body` means the parameters are in the URL query. +/// The URL is never included in errors or logs: its query may carry `client_secret`. +AccessToken requestOAuthToken(const DB::ContextPtr & context, const Poco::URI & url, const std::string & body); + +/// Calls `make_request(/* force_refresh = */ false)`. When the catalog rejects the request with +/// 401/403, the cached OAuth token is likely stale, so the request runs once more with +/// `force_refresh = true`, which makes the caller mint a fresh token first. Retries at most once. +/// `enable_refresh` is false for catalogs whose credentials cannot be renewed, e.g. a static token. +template +auto requestWithTokenRefresh(bool enable_refresh, Func && make_request) +{ + if (!enable_refresh) + return make_request(/* force_refresh = */ false); + + try + { + return make_request(/* force_refresh = */ false); + } + catch (const DB::HTTPException & e) + { + const auto status = e.getHTTPStatus(); + if (status != Poco::Net::HTTPResponse::HTTPStatus::HTTP_UNAUTHORIZED + && status != Poco::Net::HTTPResponse::HTTPStatus::HTTP_FORBIDDEN) + throw; + return make_request(/* force_refresh = */ true); + } +} + DB::ReadWriteBufferFromHTTPPtr createReadBuffer( const std::string & endpoint, DB::ContextPtr context, diff --git a/src/Databases/DataLake/ICatalog.cpp b/src/Databases/DataLake/ICatalog.cpp index 75e200177199e..bfcbcb5495aa6 100644 --- a/src/Databases/DataLake/ICatalog.cpp +++ b/src/Databases/DataLake/ICatalog.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -405,6 +406,24 @@ DB::SettingsChanges CatalogSettings::allChanged() const return changes; } +std::string_view ICatalog::getTableEngineName(const TableMetadata & table_metadata) const +{ + if (!table_metadata.isDefaultReadableTable()) + return FAKE_TABLE_ENGINE_NAME_FOR_UNREADABLE_TABLES; + + switch (getTableFormat(table_metadata)) + { + case DataLakeTableFormat::UNKNOWN: + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Table is readable, but its catalog reports no data lake format"); + case DataLakeTableFormat::DELTA: + return "DeltaLake"; + case DataLakeTableFormat::ICEBERG: + return "Iceberg"; + case DataLakeTableFormat::PAIMON: + return "Paimon"; + } +} + CatalogTables ICatalog::getTables(const TableNameFilter & filter) const { switch (filter.kind) diff --git a/src/Databases/DataLake/ICatalog.h b/src/Databases/DataLake/ICatalog.h index ef8ec304f3d81..440fd5e677702 100644 --- a/src/Databases/DataLake/ICatalog.h +++ b/src/Databases/DataLake/ICatalog.h @@ -14,6 +14,7 @@ #include #include +#include #include namespace DB @@ -50,6 +51,14 @@ class CatalogSettingsAlterValidatorFactory std::unordered_map validators; }; +enum class DataLakeTableFormat : uint8_t +{ + UNKNOWN, + DELTA, + ICEBERG, + PAIMON, +}; + struct DataLakeSpecificProperties { std::string iceberg_metadata_file_location; @@ -101,6 +110,8 @@ class TableMetadata void setTableUUID(const std::string & uuid_) { table_uuid = uuid_; } std::optional getTableUUID() const { return table_uuid; } + void setTableFormat(DataLakeTableFormat format) { table_format = format; } + DataLakeTableFormat getTableFormat() const { return table_format; } bool requiresLocation() const { return with_location; } bool requiresSchema() const { return with_schema; } @@ -157,6 +168,7 @@ class TableMetadata std::string reason_why_table_is_not_readable; std::optional table_uuid; + DataLakeTableFormat table_format = DataLakeTableFormat::UNKNOWN; bool is_default_readable_table = true; @@ -227,6 +239,22 @@ class ICatalog explicit ICatalog(const std::string & warehouse_) : warehouse(warehouse_) {} virtual DB::DatabaseDataLakeCatalogType getCatalogType() const = 0; + + /// Format of the data lake table served for `table_metadata`. Catalogs serving a single + /// format rely on the database-level table engine and never set the format per table; + /// mixed-format catalogs (e.g. `UnityV2Catalog`) set it in `getTableMetadata`. + virtual DataLakeTableFormat getTableFormat(const TableMetadata & table_metadata) const + { + return table_metadata.getTableFormat(); + } + + /// Table storage engine name for `table_metadata` ("DeltaLake", "Iceberg", "Paimon", + /// or the fake engine for unreadable tables). Throws LOGICAL_ERROR when a readable + /// table has no known format; callers handling single-format catalogs (which never set + /// the format) must check `getTableFormat() != UNKNOWN` first and keep their + /// database-level engine name otherwise. + std::string_view getTableEngineName(const TableMetadata & table_metadata) const; + virtual ~ICatalog() = default; /// Does catalog have any tables? @@ -323,7 +351,8 @@ class ICatalog /// The Glue catalog does not support such operation. virtual bool isTransactional() const { return false; } - virtual CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & /*storage_id*/) + virtual CredentialsRefreshCallback getCredentialsConfigurationCallback( + const DB::StorageID & /*storage_id*/, const TableMetadata & /*table_metadata*/) { return std::nullopt; } diff --git a/src/Databases/DataLake/RestCatalog.cpp b/src/Databases/DataLake/RestCatalog.cpp index d95b1a689433e..93e64cb7ce80f 100644 --- a/src/Databases/DataLake/RestCatalog.cpp +++ b/src/Databases/DataLake/RestCatalog.cpp @@ -144,6 +144,18 @@ static constexpr auto NAMESPACES_ENDPOINT = "namespaces"; /// so that they ask for a fresh one on the next request. static constexpr auto UNKNOWN_EXPIRATION_TOKEN_LIFETIME = std::chrono::minutes(1); +DB::HTTPHeaderEntry parseAuthHeader(const std::string & auth_header) +{ + /// Parse a string of format "Authorization: " + /// into a key-value header "Authorization", " " + + auto pos = auth_header.find(':'); + if (pos == std::string::npos) + throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Unexpected format of auth header"); + + return DB::HTTPHeaderEntry(auth_header.substr(0, pos), auth_header.substr(pos + 1)); +} + namespace { @@ -176,18 +188,6 @@ std::pair parseCatalogCredential(const std::string & c return std::pair(client_id, client_secret); } -DB::HTTPHeaderEntry parseAuthHeader(const std::string & auth_header) -{ - /// Parse a string of format "Authorization: " - /// into a key-value header "Authorization", " " - - auto pos = auth_header.find(':'); - if (pos == std::string::npos) - throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Unexpected format of auth header"); - - return DB::HTTPHeaderEntry(auth_header.substr(0, pos), auth_header.substr(pos + 1)); -} - std::string correctAPIURI(const std::string & uri) { if (uri.ends_with("v1")) @@ -1071,8 +1071,6 @@ AccessToken RestCatalog::retrieveAccessToken(const std::string & client_id, cons /// https://github.com/apache/iceberg/blob/918f81f3c3f498f46afcea17c1ac9cdc6913cb5c/open-api/rest-catalog-open-api.yaml#L183C82-L183C99 Poco::URI url; - DB::ReadWriteBufferFromHTTP::OutStreamCallback out_stream_callback; - size_t body_size = 0; String body; if (oauth_server_uri.empty() && !oauth_server_use_request_body) @@ -1105,11 +1103,6 @@ AccessToken RestCatalog::retrieveAccessToken(const std::string & client_id, cons Poco::URI::encode(client_id, client_id, encoded_client_id); body += "&client_id=" + encoded_client_id; } - body_size = body.size(); - out_stream_callback = [&](std::ostream & os) - { - os << body; - }; if (oauth_server_uri.empty()) url = Poco::URI(base_url / oauth_tokens_endpoint); @@ -1117,58 +1110,7 @@ AccessToken RestCatalog::retrieveAccessToken(const std::string & client_id, cons url = Poco::URI(oauth_server_uri); } - const auto & context = getContext(); - context->getRemoteHostFilter().checkHostAndPort(url.getHost(), std::to_string(url.getPort())); - auto timeouts = DB::ConnectionTimeouts::getHTTPTimeouts(context->getSettingsRef(), context->getServerSettings()); - auto session = makeHTTPSession(DB::HTTPConnectionGroupType::HTTP, url, timeouts, {}); - - Poco::Net::HTTPRequest request(Poco::Net::HTTPRequest::HTTP_POST, url.getPathAndQuery(), - Poco::Net::HTTPMessage::HTTP_1_1); - request.setContentType("application/x-www-form-urlencoded"); - request.setContentLength(body_size); - request.set("Accept", "application/json"); - - std::ostream & os = session->sendRequest(request); - /// The query-parameters flavor of the request has no body. - if (out_stream_callback) - out_stream_callback(os); - - Poco::Net::HTTPResponse response; - std::istream & rs = session->receiveResponse(response); - - std::string json_str; - Poco::StreamCopier::copyToString(rs, json_str); - - /// The body of a failed response is an OAuth error object, safe to show. - /// The URL is omitted: its query string can carry `client_secret`. - if (response.getStatus() != Poco::Net::HTTPResponse::HTTP_OK) - throw DB::Exception( - DB::ErrorCodes::DATALAKE_DATABASE_ERROR, - "OAuth token request failed with status {} ({}): {}", - static_cast(response.getStatus()), response.getReason(), json_str); - - Poco::JSON::Parser parser; - Poco::Dynamic::Var res_json = parser.parse(json_str); - const Poco::JSON::Object::Ptr & object = res_json.extract(); - - if (!object->has("access_token")) - throw DB::Exception( - DB::ErrorCodes::DATALAKE_DATABASE_ERROR, - "OAuth token response has no `access_token` field: {}", - json_str); - - AccessToken token; - token.token = object->getValue("access_token"); - - if (object->has("expires_in")) - { - Int64 expires_in = object->getValue("expires_in"); - /// Use 90% of the token lifetime as the validity window so that short-lived tokens - /// (e.g. expires_in=300) still get a sensible buffer instead of going non-positive. - token.expires_at = std::chrono::system_clock::now() + std::chrono::seconds(expires_in * 9 / 10); - } - - return token; + return requestOAuthToken(getContext(), url, body); } AccessToken OneLakeCatalog::retrieveAccessTokenViaRefreshToken(const CatalogState & catalog_state) const @@ -2738,7 +2680,7 @@ void RestCatalog::cacheCredentials( = VendedStorageCredentials{parsed.credentials, parsed.endpoint, refresh_after, parsed.table_uuid, parsed.location}; } -ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCallback(const DB::StorageID & storage_id) +ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCallback(const DB::StorageID & storage_id, const TableMetadata & /*table_metadata*/) { return [this, storage_id] () -> std::shared_ptr { diff --git a/src/Databases/DataLake/RestCatalog.h b/src/Databases/DataLake/RestCatalog.h index e467362e24ea7..7b5a0370066f2 100644 --- a/src/Databases/DataLake/RestCatalog.h +++ b/src/Databases/DataLake/RestCatalog.h @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -25,18 +26,8 @@ class ReadBuffer; namespace DataLake { -struct AccessToken -{ - std::string token; - std::optional expires_at; - - bool isExpired() const - { - if (!expires_at.has_value()) - return false; - return std::chrono::system_clock::now() >= expires_at.value(); - } -}; +/// Parses "Name: value" into a header entry. +DB::HTTPHeaderEntry parseAuthHeader(const std::string & auth_header); struct VendedStorageCredentials { @@ -115,7 +106,7 @@ class RestCatalog : public ICatalog, public DB::WithContext void dropTable(const String & namespace_name, const String & table_name, bool purge, bool if_exists) const override; - ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & storage_id) override; + ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & storage_id, const TableMetadata & table_metadata) override; void setVendedCredentialsCacheTTL(std::chrono::seconds ttl) override; diff --git a/src/Databases/DataLake/UnityCatalog.cpp b/src/Databases/DataLake/UnityCatalog.cpp index de1e4c5d967db..bfa5badb4878d 100644 --- a/src/Databases/DataLake/UnityCatalog.cpp +++ b/src/Databases/DataLake/UnityCatalog.cpp @@ -537,7 +537,7 @@ bool UnityCatalog::isNamespaceAllowed(const std::string & namespace_) const return allowed_namespaces.contains("*") || allowed_namespaces.contains(namespace_); } -ICatalog::CredentialsRefreshCallback UnityCatalog::getCredentialsConfigurationCallback(const DB::StorageID & table_id) +ICatalog::CredentialsRefreshCallback UnityCatalog::getCredentialsConfigurationCallback(const DB::StorageID & table_id, const TableMetadata & /*table_metadata*/) { if (!table_id.hasUUID()) throw DB::Exception( diff --git a/src/Databases/DataLake/UnityCatalog.h b/src/Databases/DataLake/UnityCatalog.h index 7ad9684d7a5ba..f321a99835d5f 100644 --- a/src/Databases/DataLake/UnityCatalog.h +++ b/src/Databases/DataLake/UnityCatalog.h @@ -84,7 +84,7 @@ class UnityCatalog final : public ICatalog, private DB::WithContext const std::string & table_name, TableMetadata & result) const; - ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & table_id) override; + ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & table_id, const TableMetadata & table_metadata) override; }; } diff --git a/src/Databases/DataLake/UnityCatalogUtils.cpp b/src/Databases/DataLake/UnityCatalogUtils.cpp new file mode 100644 index 0000000000000..2395b2f35c281 --- /dev/null +++ b/src/Databases/DataLake/UnityCatalogUtils.cpp @@ -0,0 +1,121 @@ +#include + +#if USE_PARQUET + +#include +#include +#include +#include + +namespace DB::ErrorCodes +{ + extern const int BAD_ARGUMENTS; +} + +namespace DataLake +{ + +/// Delta primitive type name (see `DeltaLakeMetadata::getSimpleTypeByName`) -> Unity `ColumnTypeName`. +static std::string deltaPrimitiveToUnityTypeName(const std::string & delta_type) +{ + if (delta_type == "boolean") return "BOOLEAN"; + if (delta_type == "byte") return "BYTE"; + if (delta_type == "short") return "SHORT"; + if (delta_type == "integer") return "INT"; + if (delta_type == "long") return "LONG"; + if (delta_type == "float") return "FLOAT"; + if (delta_type == "double") return "DOUBLE"; + if (delta_type == "date") return "DATE"; + if (delta_type == "timestamp") return "TIMESTAMP"; + if (delta_type == "timestamp_ntz") return "TIMESTAMP_NTZ"; + if (delta_type == "string") return "STRING"; + if (delta_type == "binary") return "BINARY"; + if (delta_type.starts_with("decimal(")) return "DECIMAL"; + throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Cannot map Delta type `{}` to a Unity column type", delta_type); +} + +Poco::JSON::Array::Ptr buildUnityColumnsFromDeltaSchema(const Poco::JSON::Array::Ptr & fields) +{ + Poco::JSON::Array::Ptr columns = new Poco::JSON::Array; + for (size_t i = 0; i < fields->size(); ++i) + { + auto field = fields->getObject(static_cast(i)); + const String name = field->getValue("name"); + const bool nullable = field->getValue("nullable"); + auto type_var = field->get("type"); + + Poco::JSON::Object::Ptr column = new Poco::JSON::Object; + column->set("name", name); + column->set("nullable", nullable); + column->set("position", static_cast(i)); + + int precision = 0; + int scale = 0; + String type_name; + String type_text; + String type_json; + + if (type_var.isString()) + { + const String & delta_type = type_var.extract(); + type_text = delta_type; + type_json = '"' + delta_type + '"'; + type_name = deltaPrimitiveToUnityTypeName(delta_type); + if (type_name == "DECIMAL") + { + const auto decimal_type = DB::DeltaLakeMetadata::getSimpleTypeByName(delta_type); + precision = static_cast(DB::getDecimalPrecision(*decimal_type)); + scale = static_cast(DB::getDecimalScale(*decimal_type)); + } + } + else + { + const auto & descriptor = type_var.extract(); + const String kind = descriptor->getValue("type"); + if (kind == "array") type_name = "ARRAY"; + else if (kind == "map") type_name = "MAP"; + else if (kind == "struct") type_name = "STRUCT"; + else + throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Unexpected complex Delta type `{}`", kind); + type_text = kind; + + /// Wrap so the read path's `getFieldType(parsed, "type")` sees the descriptor under `type`. + Poco::JSON::Object::Ptr wrapper = new Poco::JSON::Object; + wrapper->set("type", descriptor); + std::ostringstream oss; // STYLE_CHECK_ALLOW_STD_STRING_STREAM + wrapper->stringify(oss); + type_json = oss.str(); + } + + column->set("type_name", type_name); + column->set("type_text", type_text); + column->set("type_json", type_json); + column->set("type_precision", precision); + column->set("type_scale", scale); + columns->add(column); + } + return columns; +} + +Poco::JSON::Object::Ptr buildUnityCreateTableBody( + const String & catalog_name, + const String & schema_name, + const String & table_name, + const String & storage_location, + Poco::JSON::Array::Ptr columns) +{ + Poco::JSON::Object::Ptr body = new Poco::JSON::Object; + body->set("name", table_name); + body->set("catalog_name", catalog_name); + body->set("schema_name", schema_name); + body->set("table_type", "EXTERNAL"); + body->set("data_source_format", "DELTA"); + body->set("storage_location", storage_location); + body->set("columns", columns); + body->set("properties", Poco::JSON::Object::Ptr(new Poco::JSON::Object)); + return body; +} + +} + +#endif diff --git a/src/Databases/DataLake/UnityCatalogUtils.h b/src/Databases/DataLake/UnityCatalogUtils.h new file mode 100644 index 0000000000000..8ee523dd06204 --- /dev/null +++ b/src/Databases/DataLake/UnityCatalogUtils.h @@ -0,0 +1,26 @@ +#pragma once +#include "config.h" + +#if USE_PARQUET + +#include +#include +#include + +namespace DataLake +{ + +/// Delta schema fields -> Unity `ColumnInfo` array; `type_json` must match what the read path parses back. +Poco::JSON::Array::Ptr buildUnityColumnsFromDeltaSchema(const Poco::JSON::Array::Ptr & fields); + +/// Body of `POST /tables` registering an external Delta table. +Poco::JSON::Object::Ptr buildUnityCreateTableBody( + const String & catalog_name, + const String & schema_name, + const String & table_name, + const String & storage_location, + Poco::JSON::Array::Ptr columns); + +} + +#endif diff --git a/src/Databases/DataLake/UnityV2Catalog.cpp b/src/Databases/DataLake/UnityV2Catalog.cpp new file mode 100644 index 0000000000000..00a89592f6927 --- /dev/null +++ b/src/Databases/DataLake/UnityV2Catalog.cpp @@ -0,0 +1,915 @@ +#include + +#if USE_AVRO && USE_PARQUET + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace DB::ErrorCodes +{ + extern const int DATALAKE_DATABASE_ERROR; + extern const int LOGICAL_ERROR; + extern const int BAD_ARGUMENTS; +} + +namespace DB::DatabaseDataLakeSetting +{ + extern const DatabaseDataLakeSettingsBool use_unity_catalog_v2; +} + +namespace +{ + +bool hasValueAndItsNotNone(const std::string & value, const Poco::JSON::Object::Ptr & object) +{ + return object->has(value) && !object->isNull(value) && !object->get(value).isEmpty(); +} + +/// `ALTER DATABASE ... MODIFY SETTING` on a Unity catalog may only switch the implementation. +[[maybe_unused]] const bool unity_settings_alter_validator_registered = [] +{ + DataLake::CatalogSettingsAlterValidatorFactory::instance().registerValidator( + DB::DatabaseDataLakeCatalogType::UNITY, + [](const DB::DatabaseDataLakeSettings &, const DB::SettingsChanges & changes) + { + const auto & alterable_setting + = DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::use_unity_catalog_v2); + for (const auto & change : changes) + { + if (change.name != alterable_setting) + throw DB::Exception( + DB::ErrorCodes::BAD_ARGUMENTS, + "Setting `{}` cannot be altered for a Unity catalog (alterable settings are: {})", + change.name, + alterable_setting); + } + }); + return true; +}(); + +} + +namespace DataLake +{ + +static const auto SCHEMAS_ENDPOINT = "schemas"; +static const auto TABLES_ENDPOINT = "tables"; +static const auto TEMPORARY_CREDENTIALS_ENDPOINT = "temporary-table-credentials"; + +/// Formats handled by the Delta reader. UniForm (`DELTA_UNIFORM_ICEBERG`) can be read as either +/// Delta or Iceberg, and we choose to read it as Delta because the Delta log is the source of truth. +static const std::unordered_set DELTA_READER_FORMATS = {"DELTA", "DELTA_UNIFORM_ICEBERG"}; +/// Formats handled by the Iceberg reader (via the Iceberg REST catalog endpoint). +static const std::unordered_set ICEBERG_READER_FORMATS = {"ICEBERG"}; + +/// Other table types (views, materialized views, streaming tables, foreign tables, shallow clones) +/// do not support direct reads of their storage. +static const std::unordered_set READABLE_TABLE_TYPES = {"MANAGED", "EXTERNAL"}; + +/// An absent `table_type` does not reject the table; the format checks handle that case. +static bool hasReadableTableType(const Poco::JSON::Object::Ptr & table_json) +{ + return !hasValueAndItsNotNone("table_type", table_json) + || READABLE_TABLE_TYPES.contains(table_json->get("table_type").extract()); +} + +/// Databricks reports managed Iceberg tables with `data_source_format` == `DELTA`, +/// so only `securable_kind` tells them apart from plain Delta tables. +static const std::unordered_set ICEBERG_SECURABLE_KINDS = { + "TABLE_DELTA_ICEBERG_MANAGED", + "TABLE_DELTA_ICEBERG_EXTERNAL", +}; + +static bool hasIcebergSecurableKind(const Poco::JSON::Object::Ptr & table_json) +{ + return hasValueAndItsNotNone("securable_kind", table_json) + && ICEBERG_SECURABLE_KINDS.contains(table_json->get("securable_kind").extract()); +} + +/// Backwards compatibility with the pre-unified Unity catalog, which accepted a table by +/// `securable_kind` alone. Delta tables are documented to always carry `data_source_format`, +/// so this likely never fires, but it keeps the shipped acceptance rule for records without the format. +static bool hasCompatDeltaSecurableKind(const Poco::JSON::Object::Ptr & table_json) +{ + static const std::unordered_set compat_delta_kinds = {"TABLE_DELTA", "TABLE_DELTA_EXTERNAL"}; + return hasValueAndItsNotNone("securable_kind", table_json) + && compat_delta_kinds.contains(table_json->get("securable_kind").extract()); +} + +struct UnityV2CatalogFullSchemaName +{ + std::string catalog_name; + std::string schema_name; +}; + +static UnityV2CatalogFullSchemaName parseFullSchemaName(const std::string & full_name) +{ + auto first_dot = full_name.find('.'); + auto catalog_name = full_name.substr(0, first_dot); + auto schema = full_name.substr(first_dot + 1); + return UnityV2CatalogFullSchemaName{.catalog_name = catalog_name, .schema_name = schema}; +} + +UnityV2Catalog::UnityV2Catalog( + const std::string & catalog_, + const std::string & base_url_, + const std::string & catalog_credential_, + const std::string & auth_scope_, + const std::string & auth_header_, + const std::string & oauth_server_uri_, + DB::ContextPtr context_) + : ICatalog(catalog_) + , DB::WithContext(context_) + , base_url_str(base_url_) + , base_url(base_url_) + , log(getLogger("UnityV2Catalog(" + catalog_ + ")")) + , auth_scope(auth_scope_) + , oauth_server_uri(oauth_server_uri_) +{ + maybeSetAuthHeader(catalog_credential_, auth_header_); + if (auth_header) + return; + + auto colon_pos = catalog_credential_.find(':'); + if (colon_pos != std::string::npos) + { + use_oauth = true; + client_id = catalog_credential_.substr(0, colon_pos); + client_secret = catalog_credential_.substr(colon_pos + 1); + } + else + { + use_oauth = false; + access_token = AccessToken{.token = catalog_credential_, .expires_at = std::nullopt}; + } +} + +UnityV2Catalog::~UnityV2Catalog() = default; + +void UnityV2Catalog::maybeSetAuthHeader(const std::string & catalog_credential_, const std::string & auth_header_) +{ + if (!catalog_credential_.empty() || auth_header_.empty()) + return; + + auth_header = parseAuthHeader(auth_header_); + /// CREATE checks `http_forbid_headers`, ATTACH does not. + DB::HTTPHeaderEntries header_to_check{*auth_header}; + getContext()->getGlobalContext()->getHTTPHeaderFilter().checkAndNormalizeHeaders(header_to_check); +} + +std::string UnityV2Catalog::getOAuthRequestParams() const +{ + String encoded_auth_scope; + String encoded_client_id; + String encoded_client_secret; + Poco::URI::encode(auth_scope, auth_scope, encoded_auth_scope); + Poco::URI::encode(client_id, client_id, encoded_client_id); + Poco::URI::encode(client_secret, client_secret, encoded_client_secret); + + return fmt::format( + "grant_type=client_credentials&scope={}&client_id={}&client_secret={}", + encoded_auth_scope, encoded_client_id, encoded_client_secret); +} + +AccessToken UnityV2Catalog::retrieveAccessToken() const +{ + Poco::URI url; + if (oauth_server_uri.empty()) + { + url = Poco::URI(base_url_str); + url.setPathEtc("/oidc/v1/token"); + } + else + url = Poco::URI(oauth_server_uri); + + return requestOAuthToken(getContext(), url, getOAuthRequestParams()); +} + +void UnityV2Catalog::ensureBearerToken(bool force_refresh) const +{ + if (access_token.has_value()) + { + if (!use_oauth || (!force_refresh && !access_token->isExpired())) + return; + + LOG_DEBUG(log, "Refreshing bearer token via OAuth ({})", force_refresh ? "rejected by the catalog" : "expired"); + access_token.reset(); + /// `iceberg_rest_catalog` keeps the bearer token inside its auth header, so it must be recreated with the new token. + iceberg_rest_catalog.reset(); + } + + if (!use_oauth) + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "No bearer token and not using OAuth"); + + access_token = retrieveAccessToken(); +} + +std::string UnityV2Catalog::getBearerToken(bool force_refresh) const +{ + std::lock_guard lock(token_mutex); + ensureBearerToken(force_refresh); + return access_token->token; +} + +DB::HTTPHeaderEntries UnityV2Catalog::getAuthHeaders(bool force_refresh) const +{ + if (auth_header) + return {*auth_header}; + + auto token = getBearerToken(force_refresh); + /// Empty token: anonymous deployment. + if (token.empty()) + return {}; + + DB::HTTPHeaderEntries headers{{"Authorization", "Bearer " + token}}; + getContext()->getGlobalContext()->getHTTPHeaderFilter().checkAndNormalizeHeaders(headers); + return headers; +} + +template +auto UnityV2Catalog::requestWithRetry(Func && make_request) const +{ + return requestWithTokenRefresh(/* enable_refresh = */ use_oauth, std::forward(make_request)); +} + +std::pair UnityV2Catalog::getJSONRequest( + const std::string & route, + const Poco::URI::QueryParameters & params) const +{ + return requestWithRetry([&](bool force_refresh) + { + return makeHTTPRequestAndReadJSON( + base_url / route, getContext(), /* bearer_token = */ "", params, getAuthHeaders(force_refresh)); + }); +} + +std::pair UnityV2Catalog::postJSONRequest( + const std::string & route, + std::function out_stream_callback) const +{ + /// `out_stream_callback` is copied, not moved: the retry has to send the same body again. + return requestWithRetry([&](bool force_refresh) + { + auto headers = getAuthHeaders(force_refresh); + /// Some Unity servers answer HTTP 500 to a POST without an explicit `Content-Type`. + headers.emplace_back("Content-Type", "application/json"); + return makeHTTPRequestAndReadJSON( + base_url / route, getContext(), /* bearer_token = */ "", {}, headers, + Poco::Net::HTTPRequest::HTTP_POST, out_stream_callback); + }); +} + +DataLakeTableFormat UnityV2Catalog::detectTableFormat(const Poco::JSON::Object::Ptr & table_json) const +{ + /// Checked before the format, because managed Iceberg tables report `data_source_format` == `DELTA`. + if (hasIcebergSecurableKind(table_json)) + return DataLakeTableFormat::ICEBERG; + + if (!hasValueAndItsNotNone("data_source_format", table_json)) + { + if (hasCompatDeltaSecurableKind(table_json)) + return DataLakeTableFormat::DELTA; + + LOG_DEBUG(log, "Table JSON has no data_source_format"); + return DataLakeTableFormat::UNKNOWN; + } + + auto format = table_json->get("data_source_format").extract(); + if (DELTA_READER_FORMATS.contains(format)) + return DataLakeTableFormat::DELTA; + if (ICEBERG_READER_FORMATS.contains(format)) + return DataLakeTableFormat::ICEBERG; + + LOG_DEBUG(log, "Unrecognized data_source_format: '{}'", format); + return DataLakeTableFormat::UNKNOWN; +} + +bool UnityV2Catalog::empty() const +{ + auto all_schemas = getSchemas(""); + for (const auto & schema : all_schemas) + { + if (!getTablesForSchema(schema, 1).empty()) + return false; + } + return true; +} + +CatalogTables UnityV2Catalog::getTables() const +{ + CatalogTables result; + auto all_schemas = getSchemas(""); + for (const auto & schema : all_schemas) + { + auto schema_tables = getTablesForSchema(schema); + result.insert(result.end(), schema_tables.begin(), schema_tables.end()); + } + return result; +} + +DataLake::ICatalog::Namespaces UnityV2Catalog::getNamespaces() const +{ + /// Unity schemas are flat — they cannot contain nested namespaces. + return getSchemas(""); +} + +CatalogTables UnityV2Catalog::listTablesInNamespaceDirect(const std::string & namespace_name) const +{ + return getTablesForSchema(namespace_name); +} + +bool UnityV2Catalog::existsTable(const std::string & schema_name, const std::string & table_name) const +{ + auto full_table_name = fmt::format("{}.{}.{}", warehouse, schema_name, table_name); + try + { + auto json = getJSONRequest(std::filesystem::path{TABLES_ENDPOINT} / full_table_name).first; + const Poco::JSON::Object::Ptr & object = json.extract(); + return hasValueAndItsNotNone("name", object) + && object->get("name").extract() == table_name; + } + catch (const DB::HTTPException & e) + { + if (e.getHTTPStatus() != Poco::Net::HTTPResponse::HTTP_NOT_FOUND) + throw; + + /// The OSS server also maps a missing catalog or schema to 404, so probe the schema + /// to keep a misconfigured namespace an error instead of an absent table. + checkNamespaceExists(schema_name); + LOG_DEBUG(log, "Unity table {} does not exist", full_table_name); + return false; + } +} + +void UnityV2Catalog::checkNamespaceExists(const std::string & schema_name) const +{ + try + { + getJSONRequest(std::filesystem::path{SCHEMAS_ENDPOINT} / fmt::format("{}.{}", warehouse, schema_name)); + } + catch (const DB::HTTPException & e) + { + if (e.getHTTPStatus() == Poco::Net::HTTPResponse::HTTP_NOT_FOUND) + throw DB::Exception( + DB::ErrorCodes::DATALAKE_DATABASE_ERROR, + "DeltaLake catalog `{}` has no schema `{}` (or the catalog itself does not exist)", + warehouse, schema_name); + throw; + } +} + +bool UnityV2Catalog::createTable( + const String & namespace_name, + const String & table_name, + const String & table_location, + Poco::JSON::Object::Ptr metadata_content, + DB::CompressionMethod metadata_compression_method, + bool if_not_exists) const +{ + /// Iceberg metadata carries `schemas`, Delta metadata carries `fields`. + if (!metadata_content->has("fields")) + { + /// Registers an external table at the location where `v1.metadata.json` was already written, like the Delta branch below. + return requestWithRetry([&](bool force_refresh) + { + return getIcebergRestCatalog(force_refresh)->createTable( + namespace_name, table_name, table_location, metadata_content, metadata_compression_method, if_not_exists); + }); + } + + auto fields = metadata_content->getArray("fields"); + if (!fields) + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Delta schema fields are missing for Unity createTable"); + + auto body = buildUnityCreateTableBody( + warehouse, namespace_name, table_name, table_location, buildUnityColumnsFromDeltaSchema(fields)); + + LOG_DEBUG(log, "Creating table {}.{}.{} at `{}` in Unity catalog", warehouse, namespace_name, table_name, table_location); + + try + { + auto response = postJSONRequest(TABLES_ENDPOINT, [&](std::ostream & os) { body->stringify(os); }); + LOG_TEST(log, "Unity createTable response: {}", response.second); + } + catch (...) + { + throw DB::Exception( + DB::ErrorCodes::DATALAKE_DATABASE_ERROR, + "Failed to create table {}.{} in Unity catalog: {}", + namespace_name, table_name, DB::getCurrentExceptionMessage(/* with_stacktrace */ false)); + } + return true; +} + +void UnityV2Catalog::createNamespaceIfNotExists(const String & namespace_name, const String & /* location */) const +{ + checkNamespaceExists(namespace_name); +} + +bool UnityV2Catalog::updateMetadata( + const String & namespace_name, + const String & table_name, + const String & new_metadata_path, + Poco::JSON::Object::Ptr new_snapshot) const +{ + /// `false` means a commit conflict (HTTP 409) and the caller retries, so the result is passed through unchanged. + return requestWithRetry([&](bool force_refresh) + { + return getIcebergRestCatalog(force_refresh)->updateMetadata(namespace_name, table_name, new_metadata_path, new_snapshot); + }); +} + +bool UnityV2Catalog::updateSchema( + const String & namespace_name, + const String & table_name, + const String & new_metadata_path, + Poco::JSON::Object::Ptr new_schema, + Int32 previous_schema_id, + Int32 new_last_column_id, + Poco::JSON::Object::Ptr metadata) const +{ + return requestWithRetry([&](bool force_refresh) + { + return getIcebergRestCatalog(force_refresh)->updateSchema( + namespace_name, table_name, new_metadata_path, new_schema, previous_schema_id, new_last_column_id, metadata); + }); +} + +void UnityV2Catalog::getTableMetadata( + const std::string & namespace_name, + const std::string & table_name, + DB::ContextPtr context_, + TableMetadata & result) const +{ + if (!tryGetTableMetadata(namespace_name, table_name, context_, result)) + throw DB::Exception(DB::ErrorCodes::DATALAKE_DATABASE_ERROR, "No response from Unity catalog"); +} + +bool UnityV2Catalog::tryGetTableMetadata( + const std::string & schema_name, + const std::string & table_name, + DB::ContextPtr context_, + TableMetadata & result) const +{ + auto full_table_name = fmt::format("{}.{}.{}", warehouse, schema_name, table_name); + + auto json = getJSONRequest(std::filesystem::path{TABLES_ENDPOINT} / full_table_name).first; + const Poco::JSON::Object::Ptr & object = json.extract(); + + if (!hasValueAndItsNotNone("name", object) || object->get("name").extract() != table_name) + return false; + + auto table_format = detectTableFormat(object); + result.setTableFormat(table_format); + + if (!hasReadableTableType(object)) + { + /// Fall through to the Delta arm: it fills the location and schema + /// best-effort, so the table stays listed with its columns. + result.setTableIsNotReadable(fmt::format( + "Cannot read table `{}` because it has table_type '{}'. " + "Readable table types are: [{}]", + full_table_name, object->get("table_type").extract(), + fmt::join(READABLE_TABLE_TYPES, ", "))); + } + else if (table_format == DataLakeTableFormat::ICEBERG) + { + /// The Unity tables API describes the table but does not serve its Iceberg metadata; + /// Databricks exposes that only through the Iceberg REST catalog endpoint. + /// See https://docs.databricks.com/aws/en/external-access/iceberg + try + { + /// Covers `getIcebergRestCatalog`: the `RestCatalog` constructor fetches `/v1/config`. + return requestWithRetry([&](bool force_refresh) + { + return getIcebergRestCatalog(force_refresh)->tryGetTableMetadata(schema_name, table_name, context_, result); + }); + } + catch (...) + { + /// Status codes for a refused request are undocumented, so any failure marks the table unreadable. + std::string iceberg_rest_url = std::filesystem::path(base_url_str) / "iceberg-rest"; + auto reason = fmt::format( + "Cannot read Iceberg table `{}`: the Unity Iceberg REST endpoint `{}` returned an error: {}. " + "Reading managed Iceberg tables from Unity requires 'External data access' to be enabled on the metastore " + "and the `EXTERNAL USE SCHEMA` privilege for the principal; " + "check these if the error above indicates a permission problem.", + full_table_name, iceberg_rest_url, DB::getCurrentExceptionMessage(false)); + LOG_DEBUG(log, "{}", reason); + result.setTableIsNotReadable(reason); + return true; + } + } + + return tryGetDeltaTableMetadata(full_table_name, object, result); +} + +bool UnityV2Catalog::tryGetDeltaTableMetadata( + const std::string & full_table_name, + const Poco::JSON::Object::Ptr & object, + TableMetadata & result) const +{ + if (result.requiresLocation()) + { + if (hasValueAndItsNotNone("storage_location", object)) + { + result.setLocation(object->get("storage_location").extract()); + } + else + { + result.setTableIsNotReadable(fmt::format( + "Cannot read table `{}` because it doesn't have storage location. " + "It means that it's not a DeltaLake table, and it's unreadable with Unity catalog in ClickHouse", + full_table_name)); + } + } + + if (hasValueAndItsNotNone("data_source_format", object)) + { + if (!DELTA_READER_FORMATS.contains(object->get("data_source_format").extract())) + { + result.setTableIsNotReadable(fmt::format( + "Cannot read table `{}` as Delta because it has data_source_format '{}'", + full_table_name, object->get("data_source_format").extract())); + } + } + else if (!hasCompatDeltaSecurableKind(object)) + { + result.setTableIsNotReadable(fmt::format( + "Cannot read table `{}` because it has no information about data_source_format", + full_table_name)); + } + + LOG_DEBUG(log, "Processing Delta table {} is default readable {}", full_table_name, result.isDefaultReadableTable()); + + if (result.requiresSchema()) + { + DB::NamesAndTypesList schema; + try + { + auto columns_json = object->getArray("columns"); + for (size_t i = 0; i < columns_json->size(); ++i) + { + const auto column_json = columns_json->get(static_cast(i)).extract(); + std::string name = column_json->getValue("name"); + auto is_nullable = column_json->getValue("nullable"); + const auto type_json_str = column_json->get("type_json").extract(); + DB::DataTypePtr data_type; + + try + { + /// OSS Unity catalog reports a simple type as a bare quoted name, such as `"integer"`, + /// where Databricks reports the JSON object every type gets there. Both forms occur. + if (type_json_str.starts_with("\"") && type_json_str.ends_with("\"") && !type_json_str.contains('{')) + { + String type_name = type_json_str.substr(1, type_json_str.size() - 2); + auto data_type_from_str = DB::DeltaLakeMetadata::getSimpleTypeByName(type_name); + data_type = is_nullable ? makeNullable(data_type_from_str) : data_type_from_str; + } + else + { + Poco::JSON::Parser parser; + auto parsed_json_type = parser.parse(type_json_str); + data_type = DB::DeltaLakeMetadata::getFieldType( + parsed_json_type.extract(), "type", is_nullable); + } + } + catch (DB::Exception & e) + { + e.addMessage("while parsing the type of column `{}` of Delta table `{}`: {}", + name, full_table_name, type_json_str); + throw; + } + schema.push_back({name, data_type}); + } + LOG_TEST(log, "Parsed schema: {}", schema.toString()); + } + catch (...) + { + /// Tables that are not Delta can use types ClickHouse has no mapping for, such as the + /// Databricks `NULL` type. They are unreadable anyway, so keep them listed without a + /// schema instead of failing the whole catalog; `SHOW CREATE TABLE` then reports no columns. + if (result.isDefaultReadableTable()) + throw; + + LOG_DEBUG( + log, + "Cannot read table `{}` because of schema parsing exception `{}`, " + "but it is not a Delta table, so we ignore this error", + full_table_name, DB::getCurrentExceptionMessage(false)); + return true; + } + + result.setSchema(schema); + } + + /// The UUID is what `getCredentialsConfigurationCallback` uses to refresh expired credentials. + if (hasValueAndItsNotNone("table_id", object)) + result.setTableUUID(object->get("table_id").extract()); + + if (result.isDefaultReadableTable() && result.requiresCredentials()) + { + const auto storage_type = parseStorageTypeFromLocation(result.getLocation()); + if (auto credentials = getDeltaCredentials(object->get("table_id"), storage_type)) + result.setStorageCredentials(credentials); + } + + return true; +} + +std::shared_ptr UnityV2Catalog::getDeltaCredentials( + const std::string & table_id, StorageType storage_type) const +{ + LOG_DEBUG(log, "Getting credentials for table {}", table_id); + if (storage_type != StorageType::S3 && storage_type != StorageType::Azure) + return nullptr; + + Poco::JSON::Object request_body; + request_body.set("table_id", table_id); + /// TODO: Change to READ_WRITE. (Be careful to not break any existing users with READ but not READ_WRITE permissions.) + request_body.set("operation", "READ"); + + auto callback = [&request_body](std::ostream & os) { request_body.stringify(os); }; + + auto [json, _] = postJSONRequest(TEMPORARY_CREDENTIALS_ENDPOINT, callback); + const Poco::JSON::Object::Ptr & response = json.extract(); + + switch (storage_type) + { + case StorageType::S3: + return parseS3Credentials(response); + case StorageType::Azure: + return parseAzureCredentials(response); + default: + return nullptr; + } +} + +std::shared_ptr UnityV2Catalog::parseS3Credentials(const Poco::JSON::Object::Ptr & response) const +{ + if (!hasValueAndItsNotNone("aws_temp_credentials", response)) + return nullptr; + + const Poco::JSON::Object::Ptr & creds_object = response->getObject("aws_temp_credentials"); + return std::make_shared( + creds_object->get("access_key_id").extract(), + creds_object->get("secret_access_key").extract(), + creds_object->get("session_token").extract()); +} + +std::shared_ptr UnityV2Catalog::parseAzureCredentials(const Poco::JSON::Object::Ptr & response) const +{ + if (!hasValueAndItsNotNone("azure_user_delegation_sas", response)) + return nullptr; + + const Poco::JSON::Object::Ptr & creds_object = response->getObject("azure_user_delegation_sas"); + return std::make_shared(creds_object->get("sas_token").extract()); +} + +/// Only S3 refreshes credentials: `StorageAzureConfiguration::createObjectStorage` discards the callback. +ICatalog::CredentialsRefreshCallback UnityV2Catalog::getCredentialsConfigurationCallback( + const DB::StorageID & table_id, const TableMetadata & table_metadata) +{ + /// Iceberg credentials are vended by the Iceberg REST catalog. + if (table_metadata.getTableFormat() == DataLakeTableFormat::ICEBERG) + { + return [this, table_id, table_metadata]() -> std::shared_ptr + { + /// Resolved per call, because refreshing the token replaces `iceberg_rest_catalog`. + return requestWithRetry([&](bool force_refresh) -> std::shared_ptr + { + auto rest_catalog = getIcebergRestCatalog(force_refresh); + auto refresh = rest_catalog->getCredentialsConfigurationCallback(table_id, table_metadata); + return refresh ? (*refresh)() : nullptr; + }); + }; + } + + /// Delta tables refresh using table_uuid and the TEMPORARY_TABLE_CREDENTIALS endpoint. + const auto table_uuid = table_metadata.getTableUUID(); + if (!table_uuid) + throw DB::Exception( + DB::ErrorCodes::BAD_ARGUMENTS, + "Cannot build a Unity credentials refresh callback for `{}`: the catalog returned no table_id", + table_id.getNameForLogs()); + + return [this, unity_table_id = *table_uuid]() -> std::shared_ptr + { + LOG_DEBUG(log, "Update credentials in the catalog"); + + return getDeltaCredentials(unity_table_id, StorageType::S3); + }; +} + +std::shared_ptr UnityV2Catalog::getIcebergRestCatalog(bool force_refresh) const +{ + std::string token; + std::string rest_auth_header; + { + std::lock_guard lock(token_mutex); + if (iceberg_rest_catalog && !force_refresh) + return iceberg_rest_catalog; + + if (auth_header) + rest_auth_header = auth_header->name + ":" + auth_header->value; + else + { + /// On `force_refresh` this resets `iceberg_rest_catalog`, so the catalog below embeds the new token. + ensureBearerToken(force_refresh); + token = access_token->token; + /// Empty token: keep the embedded catalog anonymous too. + if (!token.empty()) + rest_auth_header = "Authorization: Bearer " + token; + } + } + + std::string iceberg_rest_url = std::filesystem::path(base_url_str) / "iceberg-rest"; + + /// Built outside `token_mutex`: the `RestCatalog` ctor fetches `/v1/config` over the network, + /// and holding the lock there would stall every Delta request waiting in `getBearerToken`. + /// + /// UnityV2Catalog owns the token and the construction of the auth_header. Omit the other oauth settings -- + /// they are only used to construct the oauth_header, so they don't need to be consumed by RestCatalog. + auto catalog = std::make_shared( + warehouse, + iceberg_rest_url, + /* catalog_credential= */ "", + /* auth_scope= */ "", + rest_auth_header, + /* oauth_server_uri= */ "", + /* oauth_server_use_request_body= */ false, + /// Unity namespaces are single-level schemas. Only matters for namespace listing, which this catalog never does. + /* namespaces= */ "*", + getContext()); + + std::lock_guard lock(token_mutex); + /// A concurrent refresh may have replaced the token meanwhile; a catalog built on the old one must not be cached. + if (auth_header || (access_token && access_token->token == token)) + iceberg_rest_catalog = catalog; + + return catalog; +} + +CatalogTables UnityV2Catalog::getTablesForSchema(const std::string & schema, size_t limit) const +{ + Poco::URI::QueryParameters params; + params.push_back({"catalog_name", warehouse}); + params.push_back({"schema_name", schema}); + params.push_back({"max_results", DB::toString(limit)}); + + CatalogTables tables; + do + { + String json_str; + Poco::Dynamic::Var json; + try + { + std::tie(json, json_str) = getJSONRequest(TABLES_ENDPOINT, params); + const Poco::JSON::Object::Ptr & object = json.extract(); + + /// A page may be empty (the "tables" field is omitted) while more pages exist, + /// so fall through to the next_page_token check. + if (hasValueAndItsNotNone("tables", object)) + { + auto tables_object = object->get("tables").extract(); + if (!tables_object) + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Cannot parse result"); + + for (size_t i = 0; i < tables_object->size(); ++i) + { + const auto current_table_json = tables_object->get(static_cast(i)).extract(); + const auto table_name = current_table_json->get("name").extract(); + auto qualified_name = fmt::format("{}.{}", schema, table_name); + + const auto table_format = detectTableFormat(current_table_json); + /// Delta needs `storage_location` from this response, while an Iceberg table + /// gets its location from the Iceberg REST catalog instead. + const bool has_location = table_format == DataLakeTableFormat::ICEBERG + || hasValueAndItsNotNone("storage_location", current_table_json); + + tables.push_back(CatalogTable{ + .name = qualified_name, + .is_readable = table_format != DataLakeTableFormat::UNKNOWN + && has_location + && hasReadableTableType(current_table_json), + }); + + if (limit && tables.size() >= limit) + break; + } + } + + if (limit && tables.size() >= limit) + break; + + if (hasValueAndItsNotNone("next_page_token", object)) + { + auto continuation_token = object->get("next_page_token").extract(); + if (continuation_token.empty()) + break; + + if (params.size() == 4) + params.pop_back(); + params.push_back({"page_token", continuation_token}); + } + else + { + break; + } + } + catch (DB::Exception & e) + { + e.addMessage("while parsing JSON: " + json_str); + throw; + } + } while (true); + + return tables; +} + +ICatalog::Namespaces UnityV2Catalog::getSchemas(const std::string & base_prefix, size_t limit) const +{ + Poco::URI::QueryParameters params; + params.push_back({"catalog_name", warehouse}); + + ICatalog::Namespaces schemas; + do + { + String json_str; + Poco::Dynamic::Var json; + try + { + std::tie(json, json_str) = getJSONRequest(SCHEMAS_ENDPOINT, params); + const Poco::JSON::Object::Ptr & object = json.extract(); + + /// A page may be empty (the "schemas" field is omitted) while more pages exist, + /// so fall through to the next_page_token check. + if (hasValueAndItsNotNone("schemas", object)) + { + auto schemas_object = object->get("schemas").extract(); + if (!schemas_object) + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Cannot parse result"); + + for (size_t i = 0; i < schemas_object->size(); ++i) + { + auto schema_info = schemas_object->get(static_cast(i)).extract(); + chassert(schema_info->get("catalog_name").extract() == warehouse); + auto schema_name = parseFullSchemaName(schema_info->get("full_name").extract()); + + if (schema_name.schema_name.starts_with(base_prefix)) + schemas.push_back(schema_name.schema_name); + + if (limit && schemas.size() >= limit) + break; + } + } + + if (limit && schemas.size() >= limit) + break; + + if (hasValueAndItsNotNone("next_page_token", object)) + { + auto continuation_token = object->get("next_page_token").extract(); + if (continuation_token.empty()) + break; + + if (params.size() == 2) + params.pop_back(); + params.push_back({"page_token", continuation_token}); + } + else + { + break; + } + } + catch (DB::Exception & e) + { + e.addMessage("while parsing JSON: " + json_str); + throw; + } + } while (true); + + return schemas; +} + +} + +#endif diff --git a/src/Databases/DataLake/UnityV2Catalog.h b/src/Databases/DataLake/UnityV2Catalog.h new file mode 100644 index 0000000000000..01b72cae41720 --- /dev/null +++ b/src/Databases/DataLake/UnityV2Catalog.h @@ -0,0 +1,176 @@ +#pragma once +#include "config.h" + +#if USE_AVRO && USE_PARQUET + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace DataLake +{ + +/// Unified Unity Catalog that supports both Delta and Iceberg tables +/// in a single database, with auto-detection of table format. +/// Supports PAT, OAuth (client_id:client_secret) and `auth_header`. +class UnityV2Catalog final : public ICatalog, private DB::WithContext +{ +public: + /// catalog_credential_ is either a bearer token or "client_id:client_secret". + UnityV2Catalog( + const std::string & catalog_, + const std::string & base_url_, + const std::string & catalog_credential_, + const std::string & auth_scope_, + const std::string & auth_header_, + const std::string & oauth_server_uri_, + DB::ContextPtr context_); + + ~UnityV2Catalog() override; + + bool empty() const override; + CatalogTables getTables() const override; + Namespaces getNamespaces() const override; + CatalogTables listTablesInNamespaceDirect(const std::string & namespace_name) const override; + bool existsTable(const std::string & schema_name, const std::string & table_name) const override; + + bool createTable( + const String & namespace_name, + const String & table_name, + const String & table_location, + Poco::JSON::Object::Ptr metadata_content, + DB::CompressionMethod metadata_compression_method, + bool if_not_exists) const override; + + /// Only checks that the schema exists. Unity schemas carry ownership and grants, so `CREATE TABLE` must not create them. + void createNamespaceIfNotExists(const String & namespace_name, const String & location) const override; + + void getTableMetadata( + const std::string & namespace_name, + const std::string & table_name, + DB::ContextPtr context_, + TableMetadata & result) const override; + + bool tryGetTableMetadata( + const std::string & schema_name, + const std::string & table_name, + DB::ContextPtr context_, + TableMetadata & result) const override; + + std::optional getStorageType() const override { return std::nullopt; } + + DB::DatabaseDataLakeCatalogType getCatalogType() const override + { + return DB::DatabaseDataLakeCatalogType::UNITY; + } + + /// Serves both Delta and Iceberg tables, so the format is detected per table. + DataLakeTableFormat getTableFormat(const TableMetadata & table_metadata) const override + { + return table_metadata.getTableFormat(); + } + + ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback( + const DB::StorageID & table_id, const TableMetadata & table_metadata) override; + + /// Iceberg tables go through the Iceberg REST endpoint, which commits atomically. + /// Without this, a write would leave metadata files in the table location before failing. + bool isTransactional() const override { return true; } + + /// Iceberg commits are forwarded to the Iceberg REST endpoint. Delta writes commit through `_delta_log` and never get here. + bool updateMetadata( + const String & namespace_name, + const String & table_name, + const String & new_metadata_path, + Poco::JSON::Object::Ptr new_snapshot) const override; + + bool updateSchema( + const String & namespace_name, + const String & table_name, + const String & new_metadata_path, + Poco::JSON::Object::Ptr new_schema, + Int32 previous_schema_id, + Int32 new_last_column_id, + Poco::JSON::Object::Ptr metadata = nullptr) const override; + +private: + const std::string base_url_str; + const std::filesystem::path base_url; + const LoggerPtr log; + + /// Auth state: a user-supplied header, or a Bearer token (static or minted via OAuth). + std::optional auth_header; + std::string client_id; + std::string client_secret; + std::string auth_scope; + std::string oauth_server_uri; + bool use_oauth = false; + + /// Guards the token and everything derived from it, because `iceberg_rest_catalog` embeds the token in its auth header. + mutable std::mutex token_mutex; + mutable std::optional access_token TSA_GUARDED_BY(token_mutex); + + /// Lazy-initialized RestCatalog for Iceberg table metadata, + /// pointing to {base_url}/iceberg-rest. + mutable std::shared_ptr iceberg_rest_catalog TSA_GUARDED_BY(token_mutex); + + /// Retries `make_request` once with a fresh token when the catalog rejects the cached one. + template + auto requestWithRetry(Func && make_request) const; + + std::pair getJSONRequest( + const std::string & route, + const Poco::URI::QueryParameters & params = {}) const; + + std::pair postJSONRequest( + const std::string & route, + std::function out_stream_callback) const; + + std::string getBearerToken(bool force_refresh = false) const; + + /// `auth_header` when set, otherwise the Bearer token. + DB::HTTPHeaderEntries getAuthHeaders(bool force_refresh) const; + + /// Uses `auth_header_` only when `catalog_credential_` is empty, as in `RestCatalog`. + void maybeSetAuthHeader(const std::string & catalog_credential_, const std::string & auth_header_); + + void checkNamespaceExists(const std::string & schema_name) const; + + /// Fetches a token from the OAuth server. + AccessToken retrieveAccessToken() const; + + /// URL-encoded client-credentials grant parameters. + std::string getOAuthRequestParams() const; + + /// `force_refresh` mints a new token even when the cached one has not expired yet. + void ensureBearerToken(bool force_refresh = false) const TSA_REQUIRES(token_mutex); + + ICatalog::Namespaces getSchemas(const std::string & base_prefix, size_t limit = 0) const; + CatalogTables getTablesForSchema(const std::string & schema, size_t limit = 0) const; + + DataLakeTableFormat detectTableFormat(const Poco::JSON::Object::Ptr & table_json) const; + + bool tryGetDeltaTableMetadata( + const std::string & full_table_name, + const Poco::JSON::Object::Ptr & table_json, + TableMetadata & result) const; + + /// Asks the catalog for temporary read credentials for Delta tables. + std::shared_ptr getDeltaCredentials(const std::string & table_id, StorageType storage_type) const; + + /// Return `nullptr` when the response carries no credentials of that kind. + std::shared_ptr parseS3Credentials(const Poco::JSON::Object::Ptr & response) const; + std::shared_ptr parseAzureCredentials(const Poco::JSON::Object::Ptr & response) const; + + /// `force_refresh` rebuilds the catalog around a freshly minted token. + std::shared_ptr getIcebergRestCatalog(bool force_refresh = false) const; +}; + +} + +#endif diff --git a/tests/integration/test_database_delta/test.py b/tests/integration/test_database_delta/test.py index 6868e0568ba28..8bee8e124666d 100644 --- a/tests/integration/test_database_delta/test.py +++ b/tests/integration/test_database_delta/test.py @@ -333,14 +333,23 @@ def execute_multiple_spark_queries(node, queries_list, retry_on_timeout=False): ) +# The new Unity implementation must match the legacy one on an all-Delta +# catalog, so every test runs with both, chosen by the `use_unity_catalog_v2` +# database setting in the CREATE DATABASE query. +USE_V2_VALUES = ["0", "1"] + +UNITY_SESSION_SETTINGS = {"allow_database_unity_catalog": "1"} + + +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) @pytest.mark.parametrize("use_delta_kernel", ["1", "0"]) -def test_embedded_database_and_tables(started_cluster, use_delta_kernel): +def test_embedded_database_and_tables(started_cluster, use_delta_kernel, use_v2): test_uuid = str(uuid.uuid4()).replace("-", "_") node1 = started_cluster.instances["node1"] node1.query(f"drop database if exists unity_test_{test_uuid}") node1.query( - f"create database unity_test_{test_uuid} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, allow_experimental_delta_kernel_rs={use_delta_kernel}", - settings={"allow_experimental_database_unity_catalog": "1"}, + f"create database unity_test_{test_uuid} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, allow_experimental_delta_kernel_rs={use_delta_kernel}, use_unity_catalog_v2={use_v2}", + settings=UNITY_SESSION_SETTINGS, ) default_tables = list( sorted( @@ -382,7 +391,8 @@ def test_embedded_database_and_tables(started_cluster, use_delta_kernel): assert data_clickhouse == data_spark -def test_check_database_unity(started_cluster): +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) +def test_check_database_unity(started_cluster, use_v2): """ Test CHECK DATABASE query on Unity Catalog with a single schema. Creates one schema with multiple tables and verifies CHECK DATABASE works correctly. @@ -418,8 +428,8 @@ def test_check_database_unity(started_cluster): # Create ClickHouse database pointing to Unity Catalog node1.query( - f"create database {db_name} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') settings warehouse = 'unity', catalog_type='unity', vended_credentials=false", - settings={"allow_database_unity_catalog": "1"}, + f"create database {db_name} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, use_unity_catalog_v2={use_v2}", + settings=UNITY_SESSION_SETTINGS, ) # Verify tables are visible @@ -453,7 +463,8 @@ def test_check_database_unity(started_cluster): "SYSTEM DISABLE FAILPOINT check_database_datalake_negative" ) -def test_multiple_schemes_tables(started_cluster): +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) +def test_multiple_schemes_tables(started_cluster, use_v2): test_uuid = str(uuid.uuid4()).replace("-", "_") node1 = started_cluster.instances["node1"] # Combine schema creation, table creation and inserts into a single @@ -468,8 +479,8 @@ def test_multiple_schemes_tables(started_cluster): execute_multiple_spark_queries(node1, queries) node1.query( - f"create database multi_schema_test{test_uuid} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') settings warehouse = 'unity', catalog_type='unity', vended_credentials=false", - settings={"allow_database_unity_catalog": "1"}, + f"create database multi_schema_test{test_uuid} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, use_unity_catalog_v2={use_v2}", + settings=UNITY_SESSION_SETTINGS, ) multi_schema_tables = list( sorted( @@ -497,8 +508,9 @@ def test_multiple_schemes_tables(started_cluster): ) +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) @pytest.mark.parametrize("use_delta_kernel", ["1", "0"]) -def test_complex_table_schema(started_cluster, use_delta_kernel): +def test_complex_table_schema(started_cluster, use_delta_kernel, use_v2): node1 = started_cluster.instances["node1"] schema_name = ( f"schema_with_complex_tables_{use_delta_kernel}_{uuid.uuid4()}".replace( @@ -519,9 +531,9 @@ def test_complex_table_schema(started_cluster, use_delta_kernel): drop database if exists complex_schema; create database complex_schema engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') -settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, allow_experimental_delta_kernel_rs={use_delta_kernel} +settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, allow_experimental_delta_kernel_rs={use_delta_kernel}, use_unity_catalog_v2={use_v2} """, - settings={"allow_database_unity_catalog": "1"}, + settings=UNITY_SESSION_SETTINGS, ) complex_schema_tables = list( @@ -559,8 +571,9 @@ def test_complex_table_schema(started_cluster, use_delta_kernel): assert node1.contains_in_log("DeltaLakeMetadata: Initializing snapshot") +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) @pytest.mark.parametrize("use_delta_kernel", ["1", "0"]) -def test_timestamp_ntz(started_cluster, use_delta_kernel): +def test_timestamp_ntz(started_cluster, use_delta_kernel, use_v2): table_name_src = f"ntz_schema_{uuid.uuid4()}".replace("-", "_") node1 = started_cluster.instances["node1"] node1.query(f"drop database if exists {table_name_src}") @@ -586,9 +599,9 @@ def test_timestamp_ntz(started_cluster, use_delta_kernel): drop database if exists {table_name}; create database {table_name_src} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') -settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, allow_experimental_delta_kernel_rs={use_delta_kernel} +settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, allow_experimental_delta_kernel_rs={use_delta_kernel}, use_unity_catalog_v2={use_v2} """, - settings={"allow_database_unity_catalog": "1"}, + settings=UNITY_SESSION_SETTINGS, ) ntz_tables = list( @@ -723,7 +736,8 @@ def get_schemas(): assert schema_name in get_schemas() -def test_used_storages_in_query_log(started_cluster): +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) +def test_used_storages_in_query_log(started_cluster, use_v2): node1 = started_cluster.instances["node1"] db_name = f"db_query_log_{uuid.uuid4()}".replace("-", "_") @@ -732,9 +746,9 @@ def test_used_storages_in_query_log(started_cluster): drop database if exists {db_name}; create database {db_name} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') -settings warehouse = 'unity', catalog_type='unity', vended_credentials=false +settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, use_unity_catalog_v2={use_v2} """, - settings={"allow_database_unity_catalog": "1"}, + settings=UNITY_SESSION_SETTINGS, ) query_id = str(uuid.uuid4()).replace("-", "") @@ -752,7 +766,8 @@ def test_used_storages_in_query_log(started_cluster): assert "DeltaLake" in result, f"Expected DeltaLake in used_storages, got {result}" -def test_snapshot_version(started_cluster): +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) +def test_snapshot_version(started_cluster, use_v2): """ Test table in delta lake catalog with CDF settings (delta_lake_snapshot_start_version, delta_lake_snapshot_end_version). @@ -811,9 +826,9 @@ def get_table_versions(): drop database if exists {db_name}; create database {db_name} engine DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') -settings warehouse = 'unity', catalog_type='unity', vended_credentials=false +settings warehouse = 'unity', catalog_type='unity', vended_credentials=false, use_unity_catalog_v2={use_v2} """, - settings={"allow_database_unity_catalog": "1"}, + settings=UNITY_SESSION_SETTINGS, ) # Validate data at version 1 @@ -961,8 +976,11 @@ def get_table_versions(): ) +@pytest.mark.parametrize("use_v2", USE_V2_VALUES) @pytest.mark.parametrize("use_delta_kernel", ["1", "0"]) -def test_varchar_char_types_via_unity_catalog(started_cluster, use_delta_kernel): +def test_varchar_char_types_via_unity_catalog( + started_cluster, use_delta_kernel, use_v2 +): """ Regression test for: Unsupported DeltaLake type: varchar(n) @@ -1001,9 +1019,9 @@ def test_varchar_char_types_via_unity_catalog(started_cluster, use_delta_kernel) CREATE DATABASE {db_name} ENGINE DataLakeCatalog('http://localhost:8080/api/2.1/unity-catalog') SETTINGS warehouse = 'unity', catalog_type = 'unity', vended_credentials = false, - allow_experimental_delta_kernel_rs = {use_delta_kernel} + allow_experimental_delta_kernel_rs = {use_delta_kernel}, use_unity_catalog_v2 = {use_v2} """, - settings={"allow_experimental_database_unity_catalog": "1"}, + settings=UNITY_SESSION_SETTINGS, ) tables = ( diff --git a/tests/integration/test_database_unity_v2/__init__.py b/tests/integration/test_database_unity_v2/__init__.py new file mode 100644 index 0000000000000..e5a0d9b4834ec --- /dev/null +++ b/tests/integration/test_database_unity_v2/__init__.py @@ -0,0 +1 @@ +#!/usr/bin/env python3 diff --git a/tests/integration/test_database_unity_v2/configs/display_secrets.xml b/tests/integration/test_database_unity_v2/configs/display_secrets.xml new file mode 100644 index 0000000000000..4f463fc8cbad2 --- /dev/null +++ b/tests/integration/test_database_unity_v2/configs/display_secrets.xml @@ -0,0 +1,3 @@ + + 1 + diff --git a/tests/integration/test_database_unity_v2/configs/user_files_root.xml b/tests/integration/test_database_unity_v2/configs/user_files_root.xml new file mode 100644 index 0000000000000..6d0617ecbe1a5 --- /dev/null +++ b/tests/integration/test_database_unity_v2/configs/user_files_root.xml @@ -0,0 +1,7 @@ + + + / + diff --git a/tests/integration/test_database_unity_v2/mock_servers/uc_proxy.py b/tests/integration/test_database_unity_v2/mock_servers/uc_proxy.py new file mode 100644 index 0000000000000..63edd9c6ad150 --- /dev/null +++ b/tests/integration/test_database_unity_v2/mock_servers/uc_proxy.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""Makes the open-source Unity Catalog server look like Databricks to +`UnityV2Catalog`, for the tables in `UNIFORM_TABLES` only: + +- serves the Iceberg REST catalog at `/iceberg-rest`, not `/iceberg`; +- reports an Iceberg `securable_kind`, which the open-source server never sends; +- writes the table location as `file:///tmp/...`, which `setLocation` requires. + +It also emulates Databricks authentication: + +- `POST /oidc/v1/token` mints OAuth tokens for `CLIENT_ID:CLIENT_SECRET`, and + requires `scope=all-apis` as Databricks service principals do; +- every proxied route requires `Authorization: Bearer `, where the token + is either `PAT_TOKEN` or a minted OAuth token, and replies 401 otherwise; +- `/control/*` endpoints let the test expire tokens. +""" +import json +import re +import sys +import urllib.error +import urllib.parse +import urllib.request +import uuid +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +UPSTREAM = "http://localhost:8080" + +# The only UniForm table in the seeded sample data. +UNIFORM_TABLES = {"marksheet_uniform"} + +ICEBERG_SECURABLE_KIND = "TABLE_DELTA_ICEBERG_EXTERNAL" + +# The lookahead makes this a no-op on an already-well-formed `file:///tmp/...`. +SINGLE_SLASH_SCHEME = re.compile(r"file:/(?=tmp/marksheet_uniform)") + +# Must match the constants in test.py. +CLIENT_ID = "test-client" +CLIENT_SECRET = "test-secret" +PAT_TOKEN = "dapi-test-pat" + +# Mutated from handler threads; `set` operations are atomic under the GIL. +VALID_TOKENS = set() +PAT_VALID = True +# Principal can use the Unity tables API but not the Iceberg REST catalog. +ICEBERG_REST_DENIED = False +ICEBERG_REST_DENIED_BODY = ( + b'{"error_code": "PERMISSION_DENIED", ' + b'"message": "External data access is not enabled on the metastore"}' +) + + +def is_authorized(header): + if not header or not header.startswith("Bearer "): + return False + token = header[len("Bearer "):] + if token == PAT_TOKEN: + return PAT_VALID + return token in VALID_TOKENS + + +def normalize_scheme(data): + return SINGLE_SLASH_SCHEME.sub("file:///", data.decode()).encode() + + +def patch_table(table): + if table.get("name") in UNIFORM_TABLES: + table["securable_kind"] = ICEBERG_SECURABLE_KIND + return table + + +def patch_tables_response(data): + """The catalog uses both the paged listing and the single-table response.""" + body = json.loads(data) + if "tables" in body: + body["tables"] = [patch_table(t) for t in body["tables"]] + elif "name" in body: + body = patch_table(body) + return json.dumps(body).encode() + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + if self.path == "/": + self._reply(200, b"OK") + return + + if self.path.startswith("/control/"): + self._handle_control() + return + + if not self._check_auth(): + return + + if ICEBERG_REST_DENIED and "/iceberg-rest/" in self.path: + self._reply(403, ICEBERG_REST_DENIED_BODY) + return + + path = self.path.replace("/iceberg-rest/", "/iceberg/") + try: + response = urllib.request.urlopen(UPSTREAM + path) + status, data = response.status, response.read() + except urllib.error.HTTPError as e: + status, data = e.code, e.read() + + if status == 200: + if "/unity-catalog/tables" in path: + data = patch_tables_response(data) + elif "/unity-catalog/iceberg/" in path: + data = normalize_scheme(data) + + self._reply(status, data) + + def do_POST(self): + if self.path == "/oidc/v1/token": + self._handle_token_request() + return + + self._reply(404, b'{"error": "unsupported POST route"}') + + def _handle_control(self): + global PAT_VALID, ICEBERG_REST_DENIED + if self.path == "/control/expire": + VALID_TOKENS.clear() + self._reply(200, b"OK") + elif self.path == "/control/revoke_pat": + PAT_VALID = False + self._reply(200, b"OK") + elif self.path == "/control/restore_pat": + PAT_VALID = True + self._reply(200, b"OK") + elif self.path == "/control/deny_iceberg_rest": + ICEBERG_REST_DENIED = True + self._reply(200, b"OK") + elif self.path == "/control/allow_iceberg_rest": + ICEBERG_REST_DENIED = False + self._reply(200, b"OK") + else: + self._reply(404, b'{"error": "unknown control route"}') + + def _read_body(self): + """ClickHouse sends request bodies chunked, which `BaseHTTPRequestHandler` + does not decode.""" + if self.headers.get("Transfer-Encoding", "").lower() != "chunked": + return self.rfile.read(int(self.headers.get("Content-Length", 0))) + + chunks = [] + while True: + size = int(self.rfile.readline().strip(), 16) + if size == 0: + self.rfile.readline() + return b"".join(chunks) + chunks.append(self.rfile.read(size)) + self.rfile.readline() + + def _handle_token_request(self): + """The RFC 6749 client-credentials grant: parameters in the request body.""" + params = urllib.parse.parse_qs(self._read_body().decode()) + + if ( + params.get("grant_type") != ["client_credentials"] + or params.get("client_id") != [CLIENT_ID] + or params.get("client_secret") != [CLIENT_SECRET] + ): + self._reply(401, b'{"error": "invalid_client"}') + return + + if params.get("scope") != ["all-apis"]: + self._reply(400, b'{"error": "invalid_scope"}') + return + + token = uuid.uuid4().hex + VALID_TOKENS.add(token) + body = { + "access_token": token, + "token_type": "bearer", + "expires_in": 3600, + } + self._reply(200, json.dumps(body).encode()) + + def _check_auth(self): + if is_authorized(self.headers.get("Authorization")): + return True + self._reply(401, b'{"error": "invalid or expired token"}') + return False + + def _reply(self, status, data): + self.send_response(status) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, fmt, *args): + print("%s %s" % (self.command, self.path), flush=True) + + +if __name__ == "__main__": + ThreadingHTTPServer(("0.0.0.0", int(sys.argv[1])), Handler).serve_forever() diff --git a/tests/integration/test_database_unity_v2/test.py b/tests/integration/test_database_unity_v2/test.py new file mode 100644 index 0000000000000..3cf901f5071ee --- /dev/null +++ b/tests/integration/test_database_unity_v2/test.py @@ -0,0 +1,651 @@ +#!/usr/bin/env python3 +"""Integration tests for the new Unity catalog implementation +(`catalog_type = 'unity'` with `use_unity_catalog_v2 = 1`). + +`UnityV2Catalog` serves Delta and Iceberg tables from one catalog, detecting +the format per table. +""" + +import json +import os +import uuid + +import pytest + +from helpers.cluster import ClickHouseCluster + +CATALOG = "unity" + +UC_PORT = 8080 +PROXY_PORT = 8090 +UC_URL = f"http://localhost:{UC_PORT}/api/2.1/unity-catalog" +PROXY_URL = f"http://localhost:{PROXY_PORT}/api/2.1/unity-catalog" + +# Must match the constants in mock_servers/uc_proxy.py. +CLIENT_ID = "test-client" +CLIENT_SECRET = "test-secret" +PAT_TOKEN = "dapi-test-pat" + +# Seeded by the docker image, all Delta. +SEEDED_TABLES = [ + "default.marksheet", + "default.marksheet_uniform", + "default.numbers", + "default.user_countries", +] +UNIFORM_TABLE = "default.marksheet_uniform" +DELTA_TABLE = "default.marksheet" + +# `marksheet_uniform` is a UniForm copy of `marksheet`, byte-identical. +SEEDED_ROW_COUNT = 15 +SEEDED_FIRST_ROW = "1\tnWYHawtqUw\t930" +SEEDED_LAST_ROW = "15\tkxUUZEUoKv\t398" + +GATE_SETTING = "allow_database_unity_catalog" +V2_SETTING = "use_unity_catalog_v2" + +UC_HOME = "/tmp/unitycatalog" +UC_LOG = UC_HOME + "/uc.log" +UC_START_TIMEOUT = 120 + + +def start_unity_catalog(node): + # Make root traversable so that non-root users can access classpath files. + node.exec_in_container(["bash", "-c", "chmod a+rx /root"], user="root") + + # Copy from /unitycatalog to /tmp/unitycatalog. + node.exec_in_container( + [ + "bash", + "-c", + 'tar -C / -cf - --exclude="*/zinc" unitycatalog | tar -C /tmp -xf -', + ] + ) + + # Call start-uc-server. + node.exec_in_container( + [ + "bash", + "-c", + f"cd {UC_HOME} && bin/start-uc-server > {UC_LOG} 2>&1 &", + ] + ) + + # Wait for server to start. + try: + node.exec_in_container( + [ + "bash", + "-c", + f"for i in $(seq 1 {UC_START_TIMEOUT}); do " + "(echo > /dev/tcp/localhost/8080) 2>/dev/null && exit 0; sleep 1; done; " + f"echo 'Unity Catalog did not start within {UC_START_TIMEOUT}s' >&2; exit 1", + ] + ) + except Exception: + # A bare port-wait timeout says nothing about why the server is absent. + print( + "Unity Catalog log:\n" + + node.exec_in_container( + ["bash", "-c", f"tail -n 50 {UC_LOG} 2>&1"], nothrow=True + ) + ) + raise + + +def link_uniform_table(node): + # `marksheet_uniform` is registered at /tmp/marksheet_uniform but is + # stored in the catalog tree. Soft link them. + table_dir = UC_HOME + "/etc/data/external/unity/default/tables/marksheet_uniform" + node.exec_in_container( + [ + "bash", + "-c", + f"ln -s {table_dir} /tmp/marksheet_uniform && " + "test -d /tmp/marksheet_uniform/metadata", + ] + ) + + +PROXY_PATH = "/tmp/uc_proxy.py" +PROXY_LOG = "/tmp/uc_proxy.log" + + +def start_proxy(node): + # Copy uc_proxy.py to container. + node.copy_file_to_container( + os.path.join(os.path.dirname(__file__), "mock_servers", "uc_proxy.py"), + PROXY_PATH, + ) + + # Start proxy. + node.exec_in_container( + [ + "bash", + "-c", + f"python3 {PROXY_PATH} {PROXY_PORT} > {PROXY_LOG} 2>&1 &", + ] + ) + + # Wait for proxy. + try: + node.exec_in_container( + [ + "bash", + "-c", + "for i in $(seq 1 30); do " + f'[ "$(curl -s http://localhost:{PROXY_PORT}/)" = OK ] && exit 0; sleep 1; done; ' + f"echo 'Proxy did not answer on port {PROXY_PORT}' >&2; exit 1", + ] + ) + except Exception: + print( + "Proxy log:\n" + + node.exec_in_container( + ["bash", "-c", f"tail -n 50 {PROXY_LOG} 2>&1"], nothrow=True + ) + ) + raise + + +@pytest.fixture(scope="module") +def started_cluster(): + try: + cluster = ClickHouseCluster(__file__) + cluster.add_instance( + "node1", + main_configs=["configs/user_files_root.xml", "configs/display_secrets.xml"], + image="altinityinfra/integration-test-with-unity-catalog", + with_installed_binary=False, + stay_alive=True, + tag=os.environ.get("DOCKER_BASE_WITH_UNITY_CATALOG_TAG", "latest"), + ) + + cluster.start() + + node = cluster.instances["node1"] + start_unity_catalog(node) + link_uniform_table(node) + start_proxy(node) + + yield cluster + + finally: + cluster.shutdown() + + +def skip_if_no_delta_kernel(node): + # The Delta kernel (Rust) is not built under Memory Sanitizer, so the DeltaLake engine is absent. + has_delta_lake = ( + int( + node.query( + "SELECT count() FROM system.table_engines WHERE name = 'DeltaLake'" + ).strip() + ) + > 0 + ) + if not has_delta_lake: + pytest.skip("Build does not support DeltaLake (Delta kernel is unavailable)") + + +def unique_name(prefix): + return f"{prefix}_{uuid.uuid4()}".replace("-", "_") + + +def create_database(node, db_name, url=UC_URL, catalog_credential=None): + credential_clause = ( + f",\n catalog_credential = '{catalog_credential}'" + if catalog_credential is not None + else "" + ) + node.query(f"DROP DATABASE IF EXISTS {db_name}") + node.query( + f""" +CREATE DATABASE {db_name} ENGINE = DataLakeCatalog('{url}') +SETTINGS warehouse = '{CATALOG}', catalog_type = 'unity', {V2_SETTING} = 1, + vended_credentials = false{credential_clause} + """, + settings={GATE_SETTING: "1"}, + ) + + +def proxy_control(node, route): + """The proxy's unauthenticated side channel for the tests.""" + return node.exec_in_container( + ["curl", "-s", f"http://localhost:{PROXY_PORT}/control/{route}"] + ) + + +def show_tables(node, db_name, pattern): + result = node.query(f"SHOW TABLES FROM {db_name} LIKE '{pattern}'").strip() + return sorted(result.split("\n")) if result else [] + + +def assert_seeded_rows(node, db_name, table): + """`marksheet` and its UniForm copy hold the same rows, whichever arm reads them.""" + rows = ( + node.query(f"SELECT * FROM {db_name}.`{table}` ORDER BY 1, 2, 3") + .strip() + .split("\n") + ) + assert len(rows) == SEEDED_ROW_COUNT + assert rows[0] == SEEDED_FIRST_ROW + assert rows[-1] == SEEDED_LAST_ROW + + +def uc_api_post(node, route, payload): + """Registers objects the seeded data does not provide, such as a CSV table.""" + script = f""" +import json, urllib.request +request = urllib.request.Request( + {UC_URL + "/" + route!r}, + data={json.dumps(payload)!r}.encode(), + method="POST", + headers={{"Content-Type": "application/json"}}, +) +print(urllib.request.urlopen(request).status) +""" + return node.exec_in_container(["python3", "-c", script]) + + +def test_experimental_gate(started_cluster): + """`CREATE DATABASE` must refuse when the Unity opt-in setting is disabled. + + Antalya enables the gate by default, so the test turns it off explicitly. + """ + node = started_cluster.instances["node1"] + db_name = unique_name("gated") + + error = node.query_and_get_error( + f""" +CREATE DATABASE {db_name} ENGINE = DataLakeCatalog('{UC_URL}') +SETTINGS warehouse = '{CATALOG}', catalog_type = 'unity', {V2_SETTING} = 1, vended_credentials = false + """, + settings={GATE_SETTING: "0"}, + ) + assert GATE_SETTING in error + + +def create_legacy_database(node, db_name): + """`CREATE` through the proxy, pinned to the legacy implementation.""" + node.query(f"DROP DATABASE IF EXISTS {db_name}") + node.query( + f""" +CREATE DATABASE {db_name} ENGINE = DataLakeCatalog('{PROXY_URL}') +SETTINGS warehouse = '{CATALOG}', catalog_type = 'unity', {V2_SETTING} = 0, + vended_credentials = false, catalog_credential = '{PAT_TOKEN}' + """, + settings={GATE_SETTING: "1"}, + ) + + +def assert_legacy_hides_iceberg(node, db_name): + """The legacy implementation does not read tables with an Iceberg `securable_kind`, + which the proxy stamps on the UniForm table, so it hides them.""" + assert UNIFORM_TABLE not in show_tables(node, db_name, "default%") + assert DELTA_TABLE in show_tables(node, db_name, "default%") + + +def test_alter_switches_implementation(started_cluster): + """An existing legacy database is migrated with `ALTER DATABASE ... MODIFY SETTING`.""" + node = started_cluster.instances["node1"] + db_name = unique_name("alter_v2") + create_legacy_database(node, db_name) + assert_legacy_hides_iceberg(node, db_name) + + node.query(f"ALTER DATABASE {db_name} MODIFY SETTING {V2_SETTING} = 1") + + assert f"{V2_SETTING} = 1" in node.query(f"SHOW CREATE DATABASE {db_name}") + assert "Iceberg" in node.query(f"SHOW CREATE TABLE {db_name}.`{UNIFORM_TABLE}`") + assert_seeded_rows(node, db_name, UNIFORM_TABLE) + + # Re-applying the current value on an initialized catalog is a no-op. + node.query(f"ALTER DATABASE {db_name} MODIFY SETTING {V2_SETTING} = 1") + assert "Iceberg" in node.query(f"SHOW CREATE TABLE {db_name}.`{UNIFORM_TABLE}`") + + # The switch survives a restart, and it can be reverted. + node.restart_clickhouse() + assert "Iceberg" in node.query(f"SHOW CREATE TABLE {db_name}.`{UNIFORM_TABLE}`") + + node.query(f"ALTER DATABASE {db_name} MODIFY SETTING {V2_SETTING} = 0") + assert V2_SETTING + " = 0" in node.query(f"SHOW CREATE DATABASE {db_name}") + assert_legacy_hides_iceberg(node, db_name) + + # Only the implementation switch may be altered. + error = node.query_and_get_error( + f"ALTER DATABASE {db_name} MODIFY SETTING warehouse = 'other'" + ) + assert "cannot be altered" in error + + +def test_list_and_read_delta_tables(started_cluster): + """On an all-Delta catalog the new implementation must match the legacy one.""" + node = started_cluster.instances["node1"] + db_name = unique_name("v2_delta") + create_database(node, db_name) + + assert "DeltaLake" in node.query(f"SHOW CREATE TABLE {db_name}.`{DELTA_TABLE}`") + + assert_seeded_rows(node, db_name, DELTA_TABLE) + + +def test_unreadable_table_is_hidden(started_cluster): + """An unreadable table is hidden from listings, and naming it says why.""" + node = started_cluster.instances["node1"] + schema_name = unique_name("v2_unreadable") + db_name = unique_name("v2_unreadable_db") + + uc_api_post(node, "schemas", {"name": schema_name, "catalog_name": CATALOG}) + uc_api_post( + node, + "tables", + { + "name": "csv_table", + "catalog_name": CATALOG, + "schema_name": schema_name, + "table_type": "EXTERNAL", + "data_source_format": "CSV", + "storage_location": f"file:///var/lib/clickhouse/user_files/tmp/{schema_name}/csv_table", + "columns": [ + { + "name": "id", + "type_text": "int", + "type_json": json.dumps( + { + "name": "id", + "type": "integer", + "nullable": True, + "metadata": {}, + } + ), + "type_name": "INT", + "position": 0, + "nullable": True, + } + ], + }, + ) + + create_database(node, db_name) + + assert show_tables(node, db_name, f"{schema_name}%") == [] + + error = node.query_and_get_error( + f"SELECT * FROM {db_name}.`{schema_name}.csv_table`" + ) + assert "as Delta because it has data_source_format 'CSV'" in error + + +def test_uniform_table_reads_as_delta(started_cluster): + node = started_cluster.instances["node1"] + skip_if_no_delta_kernel(node) + + db_name = unique_name("v2_uniform") + create_database(node, db_name) + + create_table = node.query(f"SHOW CREATE TABLE {db_name}.`{UNIFORM_TABLE}`") + assert "DeltaLake" in create_table + assert "Iceberg" not in create_table + + assert_seeded_rows(node, db_name, UNIFORM_TABLE) + + +def test_iceberg_table_routes_to_iceberg_arm(started_cluster): + """Databricks reports managed Iceberg as `data_source_format = DELTA`, so + `securable_kind` must win over the format.""" + node = started_cluster.instances["node1"] + proxied_db = unique_name("v2_iceberg") + direct_db = unique_name("v2_direct") + + create_database(node, proxied_db, url=PROXY_URL, catalog_credential=PAT_TOKEN) + create_database(node, direct_db) + + proxied = node.query(f"SHOW CREATE TABLE {proxied_db}.`{UNIFORM_TABLE}`") + direct = node.query(f"SHOW CREATE TABLE {direct_db}.`{UNIFORM_TABLE}`") + + assert "Iceberg" in proxied + assert "DeltaLake" in direct + + +def test_iceberg_table_is_listed_and_readable(started_cluster): + """The Iceberg arm reads through an embedded `RestCatalog`, a different + metadata path from Delta. Both arms read the same rows, so they must agree.""" + node = started_cluster.instances["node1"] + iceberg_db = unique_name("v2_iceberg_read") + create_database(node, iceberg_db, url=PROXY_URL, catalog_credential=PAT_TOKEN) + + assert UNIFORM_TABLE in show_tables(node, iceberg_db, "default%") + + described = node.query(f"DESCRIBE TABLE {iceberg_db}.`{UNIFORM_TABLE}`") + assert "id\tNullable(Int32)" in described + assert "name\tNullable(String)" in described + assert "marks\tNullable(Int32)" in described + + assert_seeded_rows(node, iceberg_db, UNIFORM_TABLE) + + +def test_mixed_formats_in_one_database(started_cluster): + """One database serving both formats, which is the point of the engine.""" + node = started_cluster.instances["node1"] + db_name = unique_name("v2_mixed") + create_database(node, db_name, url=PROXY_URL, catalog_credential=PAT_TOKEN) + + assert show_tables(node, db_name, "default%") == SEEDED_TABLES + + def used_storages(table): + query_id = str(uuid.uuid4()).replace("-", "") + node.query(f"SELECT * FROM {db_name}.`{table}` LIMIT 1", query_id=query_id) + node.query("SYSTEM FLUSH LOGS") + return node.query( + "SELECT used_storages FROM system.query_log" + f" WHERE query_id = '{query_id}' AND type = 'QueryFinish'" + ).strip() + + delta_storages = used_storages(DELTA_TABLE) + iceberg_storages = used_storages(UNIFORM_TABLE) + + assert "DeltaLake" in delta_storages + assert "Iceberg" in iceberg_storages + + +def test_pat_token_authentication(started_cluster): + node = started_cluster.instances["node1"] + db_name = unique_name("v2_pat") + + # Succeeds with correct token. + create_database(node, db_name, url=PROXY_URL, catalog_credential=PAT_TOKEN) + assert DELTA_TABLE in show_tables(node, db_name, "default%") + assert_seeded_rows(node, db_name, DELTA_TABLE) + + # Fails with incorrect token. + bad_db = unique_name("v2_pat_bad") + create_database(node, bad_db, url=PROXY_URL, catalog_credential="dapi-wrong") + error = node.query_and_get_error(f"SHOW TABLES FROM {bad_db}") + assert "401" in error + + +def test_oauth_token_refresh(started_cluster): + # Create DB with client ID and secret. + node = started_cluster.instances["node1"] + db_name = unique_name("v2_oauth") + create_database(node, db_name, PROXY_URL, f"{CLIENT_ID}:{CLIENT_SECRET}") + assert_seeded_rows(node, db_name, DELTA_TABLE) + + # Expire the token. + proxy_control(node, "expire") + + # Verify refresh. + assert_seeded_rows(node, db_name, UNIFORM_TABLE) + assert_seeded_rows(node, db_name, DELTA_TABLE) + + # Verify failure with wrong secret. + bad_db = unique_name("v2_oauth_bad") + create_database(node, bad_db, PROXY_URL, f"{CLIENT_ID}:wrong-secret") + error = node.query_and_get_error(f"SHOW TABLES FROM {bad_db}") + assert "401" in error + + +def test_static_token_expiry(started_cluster): + node = started_cluster.instances["node1"] + db_name = unique_name("v2_pat_expired") + create_database(node, db_name, url=PROXY_URL, catalog_credential=PAT_TOKEN) + assert_seeded_rows(node, db_name, DELTA_TABLE) + + try: + proxy_control(node, "revoke_pat") + + error = node.query_and_get_error(f"SHOW TABLES FROM {db_name}") + assert "401" in error + assert "LOGICAL_ERROR" not in error + + # Ensure server still starts with an expired token. + node.restart_clickhouse() + + assert db_name in node.query("SHOW DATABASES").split("\n") + assert f"{V2_SETTING} = 1" in node.query(f"SHOW CREATE DATABASE {db_name}") + + error = node.query_and_get_error(f"SHOW TABLES FROM {db_name}") + assert "401" in error + assert "LOGICAL_ERROR" not in error + finally: + proxy_control(node, "restore_pat") + + assert_seeded_rows(node, db_name, DELTA_TABLE) + + +def test_iceberg_rest_denied_keeps_database_usable(started_cluster): + """No Iceberg REST access: the table is unreadable, the database still works.""" + node = started_cluster.instances["node1"] + db_name = unique_name("v2_iceberg_rest_denied") + create_database(node, db_name, url=PROXY_URL, catalog_credential=PAT_TOKEN) + + def column_tables(): + return set( + node.query( + f"SELECT DISTINCT table FROM system.columns WHERE database = '{db_name}'", + settings={"show_data_lake_catalogs_in_system_tables": "1"}, + ) + .strip() + .split("\n") + ) + + try: + proxy_control(node, "deny_iceberg_rest") + + # Listing does not touch the REST endpoint, so the table is still shown. + assert UNIFORM_TABLE in show_tables(node, db_name, "default%") + + # Goes through getTablesIterator. + tables = column_tables() + assert DELTA_TABLE in tables + assert UNIFORM_TABLE not in tables + + error = node.query_and_get_error(f"SELECT * FROM {db_name}.`{UNIFORM_TABLE}`") + assert "EXTERNAL USE SCHEMA" in error + assert "External data access is not enabled" in error + assert "LOGICAL_ERROR" not in error + + assert_seeded_rows(node, db_name, DELTA_TABLE) + + # Not sticky. + proxy_control(node, "allow_iceberg_rest") + assert_seeded_rows(node, db_name, UNIFORM_TABLE) + assert UNIFORM_TABLE in column_tables() + + proxy_control(node, "deny_iceberg_rest") + node.query(f"DROP DATABASE {db_name}") + finally: + proxy_control(node, "allow_iceberg_rest") + + +def test_no_secrets_leaked(started_cluster): + """`catalog_credential` must not appear in `SHOW CREATE`, `system.databases`, errors, or logs.""" + node = started_cluster.instances["node1"] + secrets = [PAT_TOKEN, CLIENT_SECRET] + query_ids = [] + + def query(sql, **kwargs): + qid = uuid.uuid4().hex + query_ids.append(qid) + return node.query(sql, query_id=qid, **kwargs) + + def create(db_name, credential): + query( + f""" +CREATE DATABASE {db_name} ENGINE = DataLakeCatalog('{PROXY_URL}') +SETTINGS warehouse = '{CATALOG}', catalog_type = 'unity', {V2_SETTING} = 1, + vended_credentials = false, catalog_credential = '{credential}' + """, + settings={GATE_SETTING: "1"}, + ) + + databases = { + unique_name("v2_leak_pat"): PAT_TOKEN, + unique_name("v2_leak_oauth"): f"{CLIENT_ID}:{CLIENT_SECRET}", + } + for db_name, credential in databases.items(): + create(db_name, credential) + # Exercise the auth path so that the token exchange gets logged. + assert DELTA_TABLE in query(f"SHOW TABLES FROM {db_name}") + + show_create = query(f"SHOW CREATE DATABASE {db_name}") + assert "[HIDDEN]" in show_create + for secret in secrets: + assert secret not in show_create + + engine_full_sql = ( + f"SELECT engine_full FROM system.databases WHERE name = '{db_name}'" + ) + engine_full = query(engine_full_sql) + assert "[HIDDEN]" in engine_full + for secret in secrets: + assert secret not in engine_full + # The value is masked, not dropped. + assert credential in node.query( + engine_full_sql, + settings={"format_display_secrets_in_show_and_select": 1}, + ) + + # Authentication errors must not echo the credential. + wrong_pat = f"dapi-wrong-{uuid.uuid4().hex}" + wrong_secret = f"wrong-secret-{uuid.uuid4().hex}" + for credential, secret in [ + (wrong_pat, wrong_pat), + (f"{CLIENT_ID}:{wrong_secret}", wrong_secret), + ]: + bad_db = unique_name("v2_leak_bad") + create(bad_db, credential) + error = node.query_and_get_error(f"SHOW TABLES FROM {bad_db}") + assert "401" in error + assert secret not in error + node.query(f"DROP DATABASE {bad_db}") + + node.query("SYSTEM FLUSH LOGS system.query_log") + node.query("SYSTEM FLUSH LOGS system.text_log") + + id_list = ", ".join(f"'{qid}'" for qid in query_ids) + # All queries must be in the log, otherwise the checks below pass on empty output. + assert node.query( + f"SELECT count() FROM system.query_log WHERE query_id IN ({id_list}) AND type = 'QueryFinish'" + ).strip() == str(len(query_ids)) + logged_queries = node.query( + f"SELECT query FROM system.query_log WHERE query_id IN ({id_list})" + ) + for secret in secrets: + assert secret not in logged_queries + + text_log_rows = node.query(f""" +SELECT message, value1, value2, value3, value4, value5, value6, value7, value8, value9, value10 +FROM system.text_log +WHERE query_id IN ({id_list}) +FORMAT JSONEachRow +""").strip() + assert text_log_rows + for line in text_log_rows.split("\n"): + for val in json.loads(line).values(): + if isinstance(val, str): + for secret in secrets: + assert secret not in val + + for db_name in databases: + node.query(f"DROP DATABASE {db_name}")