SNOW-1728000 Creating Iceberg table in e2e tests

snowflakedb · Oct 24, 2024 · 3a399be · 3a399be
1 parent 204e311
commit 3a399be
Show file tree

Hide file tree

Showing 59 changed files with 228 additions and 109 deletions.
diff --git a/test/rest_request_template/iceberg_json_aws.json b/test/rest_request_template/iceberg_json_aws.json
@@ -0,0 +1,30 @@
+{
+  "name": "SNOWFLAKE_CONNECTOR_NAME",
+  "config": {
+    "connector.class": "com.snowflake.kafka.connector.SnowflakeSinkConnector",
+    "topics": "SNOWFLAKE_TEST_TOPIC",
+    "tasks.max": "1",
+    "buffer.flush.time": "10",
+    "buffer.count.records": "100",
+    "buffer.size.bytes": "5000000",
+    "snowflake.url.name": "SNOWFLAKE_HOST",
+    "snowflake.user.name": "SNOWFLAKE_USER",
+    "snowflake.private.key": "SNOWFLAKE_PRIVATE_KEY",
+    "snowflake.database.name": "SNOWFLAKE_DATABASE",
+    "snowflake.schema.name": "SNOWFLAKE_SCHEMA",
+    "snowflake.role.name": "SNOWFLAKE_ROLE",
+    "snowflake.ingestion.method": "SNOWPIPE_STREAMING",
+    "snowflake.streaming.enable.single.buffer": true,
+    "snowflake.streaming.closeChannelsInParallel.enabled": true,
+    "snowflake.enable.schematization": "false",
+    "key.converter": "org.apache.kafka.connect.storage.StringConverter",
+    "value.converter": "org.apache.kafka.connect.json.JsonConverter",
+    "value.converter.schemas.enable": "false",
+    "jmx": "true",
+    "errors.tolerance": "all",
+    "errors.log.enable": true,
+    "errors.deadletterqueue.topic.name": "DLQ_TOPIC",
+    "errors.deadletterqueue.topic.replication.factor": 1,
+    "snowflake.streaming.iceberg.enabled": true
+  }
+}
diff --git a/test/test_executor.py b/test/test_executor.py
@@ -8,6 +8,7 @@ class TestExecutor:
     def execute(self, testSuitList, driver, nameSalt, round=1):
         try:
             for test in testSuitList:
+                test.setup()
                 driver.createConnector(test.getConfigFileName(), nameSalt)
 
             driver.startConnectorWaitTime()

diff --git a/test/test_suit/base_e2e.py b/test/test_suit/base_e2e.py
@@ -0,0 +1,3 @@
+class BaseE2eTest:
+    def setup(self):
+        pass
diff --git a/test/test_suit/iceberg_json_aws.py b/test/test_suit/iceberg_json_aws.py
@@ -0,0 +1,41 @@
+import datetime
+
+from test_suit.test_utils import RetryableError, NonRetryableError
+import json
+from time import sleep
+from test_suit.base_e2e import BaseE2eTest
+
+
+class TestIcebergJsonAws(BaseE2eTest):
+    def __init__(self, driver, nameSalt: str):
+        self.driver = driver
+        self.fileName = "iceberg_json_aws"
+        self.topic = self.fileName + nameSalt
+
+    def getConfigFileName(self):
+        return self.fileName + ".json"
+
+    def setup(self):
+        self.driver.create_iceberg_table_with_content(
+            table_name=self.topic,
+            external_volume="kafka_push_e2e_volume_aws",  # volume created manually
+        )
+
+    def send(self):
+        msg = json.dumps(
+            {
+                "id": 1,
+                "body_temperature": 36.6,
+                "name": "Steve",
+                "approved_coffee_types": ["Espresso"],
+                "animals_possessed": {"dogs": True, "cats": False},
+            }
+        )
+
+    def verify(self, round):
+        res = self.driver.select_number_of_records(self.topic)
+        print("Count records in table {}={}".format(self.topic, str(res)))
+
+    def clean(self):
+        self.driver.drop_iceberg_table(self.topic)
+        return
diff --git a/test/test_suit/resilience_tests/test_kc_delete_create.py b/test/test_suit/resilience_tests/test_kc_delete_create.py
@@ -47,8 +47,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_delete_create_chaos.py b/test/test_suit/resilience_tests/test_kc_delete_create_chaos.py
@@ -50,8 +50,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_delete_resume.py b/test/test_suit/resilience_tests/test_kc_delete_resume.py
@@ -48,8 +48,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_delete_resume_chaos.py b/test/test_suit/resilience_tests/test_kc_delete_resume_chaos.py
@@ -53,8 +53,7 @@ def verify(self, round):
         # since the pressure is applied during deletion, some of the data may be ingested, so look for a range
         goalCountUpper = self.recordNum * self.expectedsends
         goalCountLower = self.recordNum * (self.expectedsends - 1)
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count between: {} - {}".format(self.topic, str(res), str(goalCountLower), str(goalCountUpper)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_pause_create.py b/test/test_suit/resilience_tests/test_kc_pause_create.py
@@ -47,8 +47,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_pause_create_chaos.py b/test/test_suit/resilience_tests/test_kc_pause_create_chaos.py
@@ -50,8 +50,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_pause_resume.py b/test/test_suit/resilience_tests/test_kc_pause_resume.py
@@ -47,8 +47,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_pause_resume_chaos.py b/test/test_suit/resilience_tests/test_kc_pause_resume_chaos.py
@@ -50,8 +50,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_recreate.py b/test/test_suit/resilience_tests/test_kc_recreate.py
@@ -47,8 +47,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_recreate_chaos.py b/test/test_suit/resilience_tests/test_kc_recreate_chaos.py
@@ -48,8 +48,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/resilience_tests/test_kc_restart.py b/test/test_suit/resilience_tests/test_kc_restart.py
@@ -53,8 +53,7 @@ def send(self):
     def verify(self, round):
         # verify record count
         goalCount = self.recordNum * self.expectedsends
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
 
         print("Count records in table {}={}. Goal record count: {}".format(self.topic, str(res), str(goalCount)))
 

diff --git a/test/test_suit/test_auto_table_creation.py b/test/test_suit/test_auto_table_creation.py
@@ -3,10 +3,11 @@
 from confluent_kafka import avro
 from confluent_kafka.schema_registry import SchemaRegistryClient
 from confluent_kafka.schema_registry import Schema
+from test_suit.base_e2e import BaseE2eTest
 
 # SR -> Schema Registry
 # Runs only in confluent test suite environment
-class TestAutoTableCreation:
+class TestAutoTableCreation(BaseE2eTest):
     def __init__(self, driver, nameSalt, schemaRegistryAddress, testSet):
         self.driver = driver
         self.fileName = "travis_correct_auto_table_creation"
@@ -118,8 +119,7 @@ def verify(self, round):
                 raise NonRetryableError("Missing column {}".format(key))
 
 
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
         if res == 0:
             raise RetryableError()
         elif res != 100:

diff --git a/test/test_suit/test_auto_table_creation_topic2table.py b/test/test_suit/test_auto_table_creation_topic2table.py
@@ -3,10 +3,11 @@
 from confluent_kafka import avro
 from confluent_kafka.schema_registry import SchemaRegistryClient
 from confluent_kafka.schema_registry import Schema
+from test_suit.base_e2e import BaseE2eTest
 
 # SR -> Schema Registry
 # Runs only in confluent test suite environment
-class TestAutoTableCreationTopic2Table:
+class TestAutoTableCreationTopic2Table(BaseE2eTest):
     def __init__(self, driver, nameSalt, schemaRegistryAddress, testSet):
         self.driver = driver
         self.fileName = "travis_correct_auto_table_creation_topic2table"

diff --git a/test/test_suit/test_avro_avro.py b/test/test_suit/test_avro_avro.py
@@ -1,7 +1,8 @@
 from test_suit.test_utils import RetryableError, NonRetryableError
+from test_suit.base_e2e import BaseE2eTest
 
 
-class TestAvroAvro:
+class TestAvroAvro(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_avro_avro"
@@ -21,8 +22,7 @@ def send(self):
         self.driver.sendBytesData(self.topic, value, key)
 
     def verify(self, round):
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
         if res == 0:
             raise RetryableError()
         elif res != 100:

diff --git a/test/test_suit/test_avrosr_avrosr.py b/test/test_suit/test_avrosr_avrosr.py
@@ -1,8 +1,9 @@
 from test_suit.test_utils import RetryableError, NonRetryableError
 from confluent_kafka import avro
+from test_suit.base_e2e import BaseE2eTest
 
 
-class TestAvrosrAvrosr:
+class TestAvrosrAvrosr(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_avrosr_avrosr"
@@ -48,8 +49,7 @@ def send(self):
             self.topic, value, self.valueSchema, key, self.keySchema)
 
     def verify(self, round):
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
         if res == 0:
             raise RetryableError()
         elif res != 100:

diff --git a/test/test_suit/test_confluent_protobuf_protobuf.py b/test/test_suit/test_confluent_protobuf_protobuf.py
@@ -3,10 +3,11 @@
 from confluent_kafka.schema_registry.protobuf import ProtobufSerializer
 from confluent_kafka.schema_registry import SchemaRegistryClient
 from confluent_kafka import SerializingProducer
+from test_suit.base_e2e import BaseE2eTest
 
 import time
 
-class TestConfluentProtobufProtobuf:
+class TestConfluentProtobufProtobuf(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_confluent_protobuf_protobuf"
@@ -50,8 +51,7 @@ def send(self):
         self.protobufProducer.flush()
 
     def verify(self, round):
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
         if res == 0:
             raise RetryableError()
         elif res != 100:

diff --git a/test/test_suit/test_json_json.py b/test/test_suit/test_json_json.py
@@ -1,8 +1,9 @@
 from test_suit.test_utils import RetryableError, NonRetryableError
 import json
+from test_suit.base_e2e import BaseE2eTest
 
 
-class TestJsonJson:
+class TestJsonJson(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_json_json"
@@ -20,8 +21,7 @@ def send(self):
         self.driver.sendBytesData(self.topic, value, key)
 
     def verify(self, round):
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
         if res == 0:
             raise RetryableError()
         elif res != 100:

diff --git a/test/test_suit/test_multiple_topic_to_one_table_snowpipe.py b/test/test_suit/test_multiple_topic_to_one_table_snowpipe.py
@@ -2,9 +2,10 @@
 from time import sleep
 from confluent_kafka import avro
 import json
+from test_suit.base_e2e import BaseE2eTest
 
 # Runs only in confluent test suite environment
-class TestMultipleTopicToOneTableSnowpipe:
+class TestMultipleTopicToOneTableSnowpipe(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_multiple_topic_to_one_table_snowpipe"

diff --git a/test/test_suit/test_multiple_topic_to_one_table_snowpipe_streaming.py b/test/test_suit/test_multiple_topic_to_one_table_snowpipe_streaming.py
@@ -2,9 +2,10 @@
 from time import sleep
 from confluent_kafka import avro
 import json
+from test_suit.base_e2e import BaseE2eTest
 
 # Runs only in confluent test suite environment
-class TestMultipleTopicToOneTableSnowpipeStreaming:
+class TestMultipleTopicToOneTableSnowpipeStreaming(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_multiple_topic_to_one_table_snowpipe_streaming"

diff --git a/test/test_suit/test_native_complex_smt.py b/test/test_suit/test_native_complex_smt.py
@@ -1,8 +1,9 @@
 from test_suit.test_utils import RetryableError, NonRetryableError
 import json
+from test_suit.base_e2e import BaseE2eTest
 
 
-class TestNativeComplexSmt:
+class TestNativeComplexSmt(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_native_complex_smt"

diff --git a/test/test_suit/test_native_string_avrosr.py b/test/test_suit/test_native_string_avrosr.py
@@ -1,8 +1,8 @@
 from test_suit.test_utils import RetryableError, NonRetryableError
 from confluent_kafka import avro
+from test_suit.base_e2e import BaseE2eTest
 
-
-class TestNativeStringAvrosr:
+class TestNativeStringAvrosr(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_native_string_avrosr"
@@ -31,8 +31,7 @@ def send(self):
         self.driver.sendAvroSRData(self.topic, value, self.valueSchema)
 
     def verify(self, round):
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
         if res == 0:
             raise RetryableError()
         elif res != 100:

diff --git a/test/test_suit/test_native_string_json_without_schema.py b/test/test_suit/test_native_string_json_without_schema.py
@@ -1,8 +1,8 @@
 from test_suit.test_utils import RetryableError, NonRetryableError
 import json
+from test_suit.base_e2e import BaseE2eTest
 
-
-class TestNativeStringJsonWithoutSchema:
+class TestNativeStringJsonWithoutSchema(BaseE2eTest):
     def __init__(self, driver, nameSalt):
         self.driver = driver
         self.fileName = "travis_correct_native_string_json_without_schema"
@@ -19,8 +19,7 @@ def send(self):
         self.driver.sendBytesData(self.topic, value)
 
     def verify(self, round):
-        res = self.driver.snowflake_conn.cursor().execute(
-            "SELECT count(*) FROM {}".format(self.topic)).fetchone()[0]
+        res = self.driver.select_number_of_records(self.topic)
         if res == 0:
             raise RetryableError()
         elif res != 100: