Cloud Fetch e2e tests

Signed-off-by: Matthew Kim <[email protected]>

Author: mattdeekay

tests/e2e/driver_tests.py

@@ -1,3 +1,4 @@
+import itertools
 from contextlib import contextmanager
 from collections import OrderedDict
 import datetime
@@ -47,6 +48,7 @@ def __init__(self, method_name):
         # If running in local mode, just use environment variables for params.
         self.arguments = os.environ if get_args_from_env else {}
         self.arraysize = 1000
+        self.buffer_size_bytes = 10485760
 
     def connection_params(self, arguments):
         params = {
@@ -79,7 +81,7 @@ def connection(self, extra_params=()):
     @contextmanager
     def cursor(self, extra_params=()):
         with self.connection(extra_params) as conn:
-            cursor = conn.cursor(arraysize=self.arraysize)
+            cursor = conn.cursor(arraysize=self.arraysize, buffer_size_bytes=self.buffer_size_bytes)
             try:
                 yield cursor
             finally:
@@ -603,14 +605,34 @@ def test_close_connection_closes_cursors(self):
             assert op_status_at_server.operationState != ttypes.TOperationState.CLOSED_STATE
 
             conn.close()
-            
+
             # When connection closes, any cursor operations should no longer exist at the server
             with self.assertRaises(thrift.Thrift.TApplicationException) as cm:
                 op_status_at_server = ars.thrift_backend._client.GetOperationStatus(status_request)
                 if hasattr(cm, "exception"):
                     assert "RESOURCE_DOES_NOT_EXIST" in cm.exception.message
 
-
+    @skipUnless(pysql_supports_arrow(), 'needs arrow support')
+    def test_cloud_fetch(self):
+        # This test can take several minutes to run
+        limits = [100000, 600000]
+        threads = [10, 25]
+        self.buffer_size_bytes = 104857600
+        self.arraysize = 100000
+        base_query = "SELECT * FROM store_sales "
+        for num_limit, num_threads, lz4_compression in itertools.product(limits, threads, [True, False]):
+            with self.subTest(num_limit=num_limit, num_threads=num_threads, lz4_compression=lz4_compression):
+                cf_result, noop_result = None, None
+                query = base_query + "LIMIT " + str(num_limit)
+                with self.cursor({"use_cloud_fetch": True, "max_download_threads": num_threads}) as cursor:
+                    cursor.execute(query)
+                    cf_result = cursor.fetchall()
+                with self.cursor({}) as cursor:
+                    cursor.execute(query)
+                    noop_result = cursor.fetchall()
+                assert len(cf_result) == len(noop_result)
+                for i in range(len(cf_result)):
+                    assert cf_result[i] == noop_result[i]
 
 
 # use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep
@@ -792,7 +814,7 @@ def perform_remove():
 
         # Clean up after ourselves
         perform_remove()
-        
+
     def test_staging_ingestion_fails_to_modify_another_staging_user(self):
         """The server should only allow modification of the staging_ingestion_user's files
         """
@@ -908,7 +930,7 @@ def generate_file_and_path_and_queries():
 
             cursor.execute(put_query1)
             cursor.execute(put_query2)
-            
+
             with pytest.raises(Error, match="Local file operations are restricted to paths within the configured staging_allowed_local_path"):
                 cursor.execute(put_query3)