Merge pull request ClickHouse#11179 from excitoon-favorites/s3globfix

alexey-milovidov · web-flow · commit d4acd1af3aae · 2020-05-26T02:45:42.000+03:00
Fixed S3 globbing which could fail in case of more than 1000 keys and some backends
diff --git a/src/Storages/StorageS3.cpp b/src/Storages/StorageS3.cpp
@@ -24,7 +24,7 @@
 #include <DataTypes/DataTypeString.h>
 
 #include <aws/s3/S3Client.h>
-#include <aws/s3/model/ListObjectsRequest.h>
+#include <aws/s3/model/ListObjectsV2Request.h>
 
 #include <Common/parseGlobs.h>
 #include <Common/quoteString.h>
@@ -228,18 +228,18 @@ Strings listFilesWithRegexpMatching(Aws::S3::S3Client & client, const S3::URI &
         return {globbed_uri.key};
     }
 
-    Aws::S3::Model::ListObjectsRequest request;
+    Aws::S3::Model::ListObjectsV2Request request;
     request.SetBucket(globbed_uri.bucket);
     request.SetPrefix(key_prefix);
 
     re2::RE2 matcher(makeRegexpPatternFromGlobs(globbed_uri.key));
     Strings result;
-    Aws::S3::Model::ListObjectsOutcome outcome;
+    Aws::S3::Model::ListObjectsV2Outcome outcome;
     int page = 0;
     do
     {
         ++page;
-        outcome = client.ListObjects(request);
+        outcome = client.ListObjectsV2(request);
         if (!outcome.IsSuccess())
         {
             throw Exception("Could not list objects in bucket " + quoteString(request.GetBucket())
@@ -256,7 +256,7 @@ Strings listFilesWithRegexpMatching(Aws::S3::S3Client & client, const S3::URI &
                 result.emplace_back(std::move(key));
         }
 
-        request.SetMarker(outcome.GetResult().GetNextMarker());
+        request.SetContinuationToken(outcome.GetResult().GetNextContinuationToken());
     }
     while (outcome.GetResult().GetIsTruncated());
 
diff --git a/tests/integration/test_storage_s3/test.py b/tests/integration/test_storage_s3/test.py
@@ -1,6 +1,7 @@
 import json
 import logging
 import random
+import threading
 
 import pytest
 
@@ -278,3 +279,31 @@ def test_wrong_s3_syntax(cluster, s3_storage_args):
 
     query = "create table test_table_s3_syntax (id UInt32) ENGINE = S3({})".format(s3_storage_args)
     assert expected_err_msg in instance.query_and_get_error(query)
+
+
+# https://en.wikipedia.org/wiki/One_Thousand_and_One_Nights
+def test_s3_glob_scheherazade(cluster):
+    bucket = cluster.minio_bucket
+    instance = cluster.instances["dummy"]  # type: ClickHouseInstance
+    table_format = "column1 UInt32, column2 UInt32, column3 UInt32"
+    max_path = ""
+    values = "(1, 1, 1)"
+    nights_per_job = 1001 // 30
+    jobs = []
+    for night in range(0, 1001, nights_per_job):
+        def add_tales(start, end):
+            for i in range(start, end):
+                path = "night_{}/tale.csv".format(i)
+                query = "insert into table function s3('http://{}:{}/{}/{}', 'CSV', '{}') values {}".format(
+                    cluster.minio_host, cluster.minio_port, bucket, path, table_format, values)
+                run_query(instance, query)
+
+        jobs.append(threading.Thread(target=add_tales, args=(night, min(night+nights_per_job, 1001))))
+        jobs[-1].start()
+
+    for job in jobs:
+        job.join()
+
+    query = "select count(), sum(column1), sum(column2), sum(column3) from s3('http://{}:{}/{}/night_*/tale.csv', 'CSV', '{}')".format(
+        cluster.minio_redirect_host, cluster.minio_redirect_port, bucket, table_format)
+    assert run_query(instance, query).splitlines() == ["1001\t1001\t1001\t1001"]
