Update test_compatibility.py::TestCompatibility::test_export (#17631)

Author: stanislav-shchetinin

ydb/tests/functional/compatibility/test_compatibility.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import boto3
+import tempfile
 import time
 import pytest
 import logging
@@ -84,6 +85,15 @@ def setup_s3():
 
         return s3_endpoint, s3_access_key, s3_secret_key, s3_bucket
 
+    def _execute_command_and_get_result(self, command):
+        with tempfile.NamedTemporaryFile(mode='w+', delete=True) as temp_file:
+            yatest.common.execute(command, wait=True, stdout=temp_file, stderr=temp_file)
+            temp_file.flush()
+            temp_file.seek(0)
+            result = json.load(temp_file)
+            self.output_f.write(str(result) + "\n")
+            return result
+
     def change_cluster_version(self, new_binary_paths):
         binary_path_before = self.config.get_binary_paths()
         versions_on_before = self.get_nodes_version()
@@ -341,28 +351,26 @@ def test_tpch1(self, store_type):
     def test_export(self):
         s3_endpoint, s3_access_key, s3_secret_key, s3_bucket = self.s3_config
 
-        session = ydb.retry_operation_sync(lambda: self.driver.table_client.session().create())
-
-        for table_num in range(1, 6):
-            table_name = f"sample_table_{table_num}"
-
-            session.execute_scheme(
-                f"create table `{table_name}` (id Uint64, payload Utf8, PRIMARY KEY(id));"
-            )
+        with ydb.SessionPool(self.driver, size=1) as pool:
+            with pool.checkout() as session:
+                for table_num in range(1, 6):
+                    table_name = f"sample_table_{table_num}"
+                    session.execute_scheme(
+                        f"create table `{table_name}` (id Uint64, payload Utf8, PRIMARY KEY(id));"
+                    )
 
-            query = f"""INSERT INTO `{table_name}` (id, payload) VALUES
-                (1, 'Payload 1 for table {table_num}'),
-                (2, 'Payload 2 for table {table_num}'),
-                (3, 'Payload 3 for table {table_num}'),
-                (4, 'Payload 4 for table {table_num}'),
-                (5, 'Payload 5 for table {table_num}');"""
-            session.transaction().execute(
-                query, commit_tx=True
-            )
+                    query = f"""INSERT INTO `{table_name}` (id, payload) VALUES
+                        (1, 'Payload 1 for table {table_num}'),
+                        (2, 'Payload 2 for table {table_num}'),
+                        (3, 'Payload 3 for table {table_num}'),
+                        (4, 'Payload 4 for table {table_num}'),
+                        (5, 'Payload 5 for table {table_num}');"""
+                    session.transaction().execute(
+                        query, commit_tx=True
+                    )
 
         export_command = [
             yatest.common.binary_path(os.getenv("YDB_CLI_BINARY")),
-            "--verbose",
             "--endpoint",
             "grpc://localhost:%d" % self.cluster.nodes[1].grpc_port,
             "--database=/Root",
@@ -377,7 +385,48 @@ def test_export(self):
             "--secret-key",
             s3_secret_key,
             "--item",
-            "src=/Root,dst=Root"
+            "src=/Root,dst=.",
+            "--format",
+            "proto-json-base64"
+        ]
+
+        result_export = self._execute_command_and_get_result(export_command)
+
+        export_id = result_export["id"]
+        status_export = result_export["status"]
+        progress_export = result_export["metadata"]["progress"]
+
+        assert status_export == "SUCCESS"
+        assert progress_export in ["PROGRESS_PREPARING", "PROGRESS_DONE"]
+
+        operation_get_command = [
+            yatest.common.binary_path(os.getenv("YDB_CLI_BINARY")),
+            "--endpoint",
+            "grpc://localhost:%d" % self.cluster.nodes[1].grpc_port,
+            "--database=/Root",
+            "operation",
+            "get",
+            "%s" % export_id,
+            "--format",
+            "proto-json-base64"
         ]
 
-        yatest.common.execute(export_command, wait=True, stdout=self.output_f, stderr=self.output_f)
+        while progress_export != "PROGRESS_DONE":
+            result_get = self._execute_command_and_get_result(operation_get_command)
+            progress_export = result_get["metadata"]["progress"]
+
+        s3_resource = boto3.resource("s3", endpoint_url=s3_endpoint, aws_access_key_id=s3_access_key, aws_secret_access_key=s3_secret_key)
+
+        keys_expected = set()
+        for table_num in range(1, 6):
+            table_name = f"sample_table_{table_num}"
+            keys_expected.add(table_name + "/data_00.csv")
+            keys_expected.add(table_name + "/metadata.json")
+            keys_expected.add(table_name + "/scheme.pb")
+
+        bucket = s3_resource.Bucket(s3_bucket)
+        keys = set()
+        for x in list(bucket.objects.all()):
+            keys.add(x.key)
+
+        assert keys_expected <= keys