Merge pull request ClickHouse#62423 from Avogar/better-exception-s3-globs-partitions

Fix exception message during writing to partitioned s3/hdfs/azure path with globs

Author: Avogar

src/Storages/HDFS/StorageHDFS.cpp

@@ -1101,6 +1101,10 @@ SinkToStoragePtr StorageHDFS::write(const ASTPtr & query, const StorageMetadataP
 
     if (is_partitioned_implementation)
     {
+        String path = current_uri.substr(current_uri.find('/', current_uri.find("//") + 2));
+        if (PartitionedSink::replaceWildcards(path, "").find_first_of("*?{") != std::string::npos)
+            throw Exception(ErrorCodes::DATABASE_ACCESS_DENIED, "URI '{}' contains globs, so the table is in readonly mode", uris.back());
+
         return std::make_shared<PartitionedHDFSSink>(
             partition_by_ast,
             current_uri,

src/Storages/StorageAzureBlob.cpp

@@ -491,6 +491,13 @@ Poco::URI StorageAzureBlob::Configuration::getConnectionURL() const
     return Poco::URI(parsed_connection_string.BlobServiceUrl.GetAbsoluteUrl());
 }
 
+bool StorageAzureBlob::Configuration::withGlobsIgnorePartitionWildcard() const
+{
+    if (!withPartitionWildcard())
+        return withGlobs();
+
+    return PartitionedSink::replaceWildcards(getPath(), "").find_first_of("*?{") != std::string::npos;
+}
 
 StorageAzureBlob::StorageAzureBlob(
     const Configuration & configuration_,
@@ -810,7 +817,7 @@ void StorageAzureBlob::read(
     size_t max_block_size,
     size_t num_streams)
 {
-    if (partition_by && configuration.withWildcard())
+    if (partition_by && configuration.withPartitionWildcard())
         throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Reading from a partitioned Azure storage is not implemented yet");
 
     auto this_ptr = std::static_pointer_cast<StorageAzureBlob>(shared_from_this());
@@ -897,13 +904,17 @@ void ReadFromAzureBlob::initializePipeline(QueryPipelineBuilder & pipeline, cons
 
 SinkToStoragePtr StorageAzureBlob::write(const ASTPtr & query, const StorageMetadataPtr & metadata_snapshot, ContextPtr local_context, bool /*async_insert*/)
 {
+    if (configuration.withGlobsIgnorePartitionWildcard())
+        throw Exception(ErrorCodes::DATABASE_ACCESS_DENIED,
+                        "AzureBlobStorage key '{}' contains globs, so the table is in readonly mode", configuration.blob_path);
+
     auto path = configuration.blobs_paths.front();
     auto sample_block = metadata_snapshot->getSampleBlock();
     auto chosen_compression_method = chooseCompressionMethod(path, configuration.compression_method);
     auto insert_query = std::dynamic_pointer_cast<ASTInsertQuery>(query);
 
     auto partition_by_ast = insert_query ? (insert_query->partition_by ? insert_query->partition_by : partition_by) : nullptr;
-    bool is_partitioned_implementation = partition_by_ast && configuration.withWildcard();
+    bool is_partitioned_implementation = partition_by_ast && configuration.withPartitionWildcard();
 
     if (is_partitioned_implementation)
     {
@@ -919,10 +930,6 @@ SinkToStoragePtr StorageAzureBlob::write(const ASTPtr & query, const StorageMeta
     }
     else
     {
-        if (configuration.withGlobs())
-            throw Exception(ErrorCodes::DATABASE_ACCESS_DENIED,
-                            "AzureBlobStorage key '{}' contains globs, so the table is in readonly mode", configuration.blob_path);
-
         if (auto new_path = checkAndGetNewFileOnInsertIfNeeded(local_context, object_storage.get(), path, configuration.blobs_paths.size()))
         {
             configuration.blobs_paths.push_back(*new_path);

src/Storages/StorageAzureBlob.h

@@ -35,12 +35,14 @@ class StorageAzureBlob : public IStorage
 
         bool withGlobs() const { return blob_path.find_first_of("*?{") != std::string::npos; }
 
-        bool withWildcard() const
+        bool withPartitionWildcard() const
         {
             static const String PARTITION_ID_WILDCARD = "{_partition_id}";
             return blobs_paths.back().find(PARTITION_ID_WILDCARD) != String::npos;
         }
 
+        bool withGlobsIgnorePartitionWildcard() const;
+
         Poco::URI getConnectionURL() const;
 
         std::string connection_url;

src/Storages/StorageS3.cpp

@@ -1283,7 +1283,7 @@ void ReadFromStorageS3Step::createIterator(const ActionsDAG::Node * predicate)
 
 void ReadFromStorageS3Step::initializePipeline(QueryPipelineBuilder & pipeline, const BuildQueryPipelineSettings &)
 {
-    if (storage.partition_by && query_configuration.withWildcard())
+    if (storage.partition_by && query_configuration.withPartitionWildcard())
         throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Reading from a partitioned S3 storage is not implemented yet");
 
     createIterator(nullptr);
@@ -1341,12 +1341,16 @@ SinkToStoragePtr StorageS3::write(const ASTPtr & query, const StorageMetadataPtr
     auto query_configuration = updateConfigurationAndGetCopy(local_context);
     auto key = query_configuration.keys.front();
 
+    if (query_configuration.withGlobsIgnorePartitionWildcard())
+        throw Exception(ErrorCodes::DATABASE_ACCESS_DENIED,
+                        "S3 key '{}' contains globs, so the table is in readonly mode", query_configuration.url.key);
+
     auto sample_block = metadata_snapshot->getSampleBlock();
     auto chosen_compression_method = chooseCompressionMethod(query_configuration.keys.back(), query_configuration.compression_method);
     auto insert_query = std::dynamic_pointer_cast<ASTInsertQuery>(query);
 
     auto partition_by_ast = insert_query ? (insert_query->partition_by ? insert_query->partition_by : partition_by) : nullptr;
-    bool is_partitioned_implementation = partition_by_ast && query_configuration.withWildcard();
+    bool is_partitioned_implementation = partition_by_ast && query_configuration.withPartitionWildcard();
 
     if (is_partitioned_implementation)
     {
@@ -1363,10 +1367,6 @@ SinkToStoragePtr StorageS3::write(const ASTPtr & query, const StorageMetadataPtr
     }
     else
     {
-        if (query_configuration.withGlobs())
-            throw Exception(ErrorCodes::DATABASE_ACCESS_DENIED,
-                            "S3 key '{}' contains globs, so the table is in readonly mode", query_configuration.url.key);
-
         if (auto new_key = checkAndGetNewFileOnInsertIfNeeded(local_context, configuration, query_configuration.keys.front(), query_configuration.keys.size()))
         {
             query_configuration.keys.push_back(*new_key);
@@ -1530,6 +1530,14 @@ void StorageS3::Configuration::connect(const ContextPtr & context)
         credentials.GetSessionToken());
 }
 
+bool StorageS3::Configuration::withGlobsIgnorePartitionWildcard() const
+{
+    if (!withPartitionWildcard())
+        return withGlobs();
+
+    return PartitionedSink::replaceWildcards(getPath(), "").find_first_of("*?{") != std::string::npos;
+}
+
 void StorageS3::processNamedCollectionResult(StorageS3::Configuration & configuration, const NamedCollection & collection)
 {
     validateNamedCollection(collection, required_configuration_keys, optional_configuration_keys);

src/Storages/StorageS3.h

@@ -274,21 +274,23 @@ class StorageS3 : public IStorage
     {
         Configuration() = default;
 
-        String getPath() const { return url.key; }
+        const String & getPath() const { return url.key; }
 
         bool update(const ContextPtr & context);
 
         void connect(const ContextPtr & context);
 
         bool withGlobs() const { return url.key.find_first_of("*?{") != std::string::npos; }
 
-        bool withWildcard() const
+        bool withPartitionWildcard() const
         {
             static const String PARTITION_ID_WILDCARD = "{_partition_id}";
             return url.bucket.find(PARTITION_ID_WILDCARD) != String::npos
                 || keys.back().find(PARTITION_ID_WILDCARD) != String::npos;
         }
 
+        bool withGlobsIgnorePartitionWildcard() const;
+
         S3::URI url;
         S3::AuthSettings auth_settings;
         S3Settings::RequestSettings request_settings;

tests/integration/test_storage_azure_blob_storage/test.py

@@ -1323,6 +1323,20 @@ def test_format_detection(cluster):
     assert result == expected_result
 
 
+def test_write_to_globbed_partitioned_path(cluster):
+    node = cluster.instances["node"]
+    storage_account_url = cluster.env_variables["AZURITE_STORAGE_ACCOUNT_URL"]
+    account_name = "devstoreaccount1"
+    account_key = "Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw=="
+    error = azure_query(
+        node,
+        f"INSERT INTO TABLE FUNCTION azureBlobStorage('{storage_account_url}', 'cont', 'test_data_*_{{_partition_id}}', '{account_name}', '{account_key}', 'CSV', 'auto', 'x UInt64') partition by 42 select 42",
+        expect_error="true",
+    )
+
+    assert "DATABASE_ACCESS_DENIED" in error
+
+
 def test_parallel_read(cluster):
     node = cluster.instances["node"]
     connection_string = cluster.env_variables["AZURITE_CONNECTION_STRING"]

tests/integration/test_storage_hdfs/test.py

@@ -1116,6 +1116,16 @@ def test_format_detection(started_cluster):
     assert expected_result == result
 
 
+def test_write_to_globbed_partitioned_path(started_cluster):
+    node = started_cluster.instances["node1"]
+
+    error = node.query_and_get_error(
+        "insert into function hdfs('hdfs://hdfs1:9000/test_data_*_{_partition_id}.csv') partition by 42 select 42"
+    )
+
+    assert "DATABASE_ACCESS_DENIED" in error
+
+
 def test_respect_object_existence_on_partitioned_write(started_cluster):
     node = started_cluster.instances["node1"]
 

tests/queries/0_stateless/03037_s3_write_to_globbed_partitioned_path.reference

@@ -0,0 +1,4 @@
+-- Tags: no-fasttest
+
+insert into function s3('http://localhost:11111/test/data_*_{_partition_id}.csv') partition by number % 3 select * from numbers(10); -- {serverError DATABASE_ACCESS_DENIED}
+