adapating put fetch get

Author: sfc-gh-pcyrek

prober/Dockerfile

@@ -43,7 +43,7 @@ ENV PATH="${PYENV_ROOT}/shims:${PYENV_ROOT}/bin:${PATH}"
 RUN git clone --depth=1 https://github.com/pyenv/pyenv.git ${PYENV_ROOT}
 
 # Build arguments for Python versions and Snowflake connector versions
-ARG MATRIX_VERSION='{"3.10.17": ["3.14.0", "3.13.2", "3.13.1"], "3.9.22": ["3.14.0", "3.13.2", "3.13.1"], "3.8.20": ["3.14.0", "3.13.2", "3.13.1"]}'
+ARG MATRIX_VERSION='{"3.13.4": ["3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"], "3.9.22": ["3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]}'
 
 
 # Install Python versions from ARG MATRIX_VERSION

prober/probes/put_fetch_get.py

@@ -1,6 +1,7 @@
 import csv
 import os
 import random
+import sys
 
 from faker import Faker
 from probes.logging_config import initialize_logger
@@ -24,21 +25,25 @@ def generate_random_data(num_records: int, file_path: str) -> str:
     Returns:
         str: File path to CSV file
     """
-    fake = Faker()
-    with open(file_path, mode="w", newline="", encoding="utf-8") as csvfile:
-        writer = csv.writer(csvfile, quoting=csv.QUOTE_ALL)
-        writer.writerow(["id", "name", "email", "address"])
-        for i in range(1, num_records + 1):
-            writer.writerow([i, fake.name(), fake.email(), fake.address()])
-    with open(file_path, newline="", encoding="utf-8") as csvfile:
-        reader = csv.reader(csvfile)
-        rows = list(reader)
-        # Subtract 1 for the header row
-        actual_records = len(rows) - 1
-        assert actual_records == num_records, logger.error(
-            f"Expected {num_records} records, but found {actual_records}."
-        )
-    return file_path
+    try:
+        fake = Faker()
+        with open(file_path, mode="w", newline="", encoding="utf-8") as csvfile:
+            writer = csv.writer(csvfile, quoting=csv.QUOTE_ALL)
+            writer.writerow(["id", "name", "email", "address"])
+            for i in range(1, num_records + 1):
+                writer.writerow([i, fake.name(), fake.email(), fake.address()])
+        with open(file_path, newline="", encoding="utf-8") as csvfile:
+            reader = csv.reader(csvfile)
+            rows = list(reader)
+            # Subtract 1 for the header row
+            actual_records = len(rows) - 1
+            assert actual_records == num_records, logger.error(
+                f"Expected {num_records} records, but found {actual_records}."
+            )
+        return file_path
+    except Exception as e:
+        logger.error(f"Error generating random data: {e}")
+        sys.exit(1)
 
 
 def create_data_table(cursor: snowflake.connector.cursor.SnowflakeCursor) -> str:
@@ -48,20 +53,26 @@ def create_data_table(cursor: snowflake.connector.cursor.SnowflakeCursor) -> str
     Returns:
         str: The name of the created table.
     """
-    table_name = random_string(7, "test_data_")
-    create_table_query = f"""
-    CREATE OR REPLACE TABLE {table_name} (
-        id INT,
-        name STRING,
-        email STRING,
-        address STRING
-    );
-    """
-    cursor.execute(create_table_query)
-    if cursor.fetchone():
-        print({"created_table": True})
-    else:
+    try:
+        table_name = random_string(7, "test_data_")
+        create_table_query = f"""
+        CREATE OR REPLACE TABLE {table_name} (
+            id INT,
+            name STRING,
+            email STRING,
+            address STRING
+        );
+        """
+        cursor.execute(create_table_query)
+        if cursor.fetchone():
+            print({"created_table": True})
+        else:
+            print({"created_table": False})
+            sys.exit(1)
+    except Exception as e:
+        logger.error(f"Error creating table: {e}")
         print({"created_table": False})
+        sys.exit(1)
     return table_name
 
 
@@ -72,15 +83,21 @@ def create_data_stage(cursor: snowflake.connector.cursor.SnowflakeCursor) -> str
     Returns:
         str: The name of the created stage.
     """
-    stage_name = random_string(7, "test_data_stage_")
-    create_stage_query = f"CREATE OR REPLACE STAGE {stage_name};"
+    try:
+        stage_name = random_string(7, "test_data_stage_")
+        create_stage_query = f"CREATE OR REPLACE STAGE {stage_name};"
 
-    cursor.execute(create_stage_query)
-    if cursor.fetchone():
-        print({"created_stage": True})
-    else:
+        cursor.execute(create_stage_query)
+        if cursor.fetchone():
+            print({"created_stage": True})
+        else:
+            print({"created_stage": False})
+            sys.exit(1)
+        return stage_name
+    except Exception as e:
         print({"created_stage": False})
-    return stage_name
+        logger.error(f"Error creating stage: {e}")
+        sys.exit(1)
 
 
 def copy_into_table_from_stage(
@@ -94,18 +111,24 @@ def copy_into_table_from_stage(
         stage_name (str): The name of the stage from which data will be copied.
         cur (snowflake.connector.cursor.SnowflakeCursor): The cursor to execute the SQL command.
     """
-    cur.execute(
-        f"""
-                    COPY INTO {table_name}
-                    FROM @{stage_name}
-                    FILE_FORMAT = (TYPE = CSV FIELD_OPTIONALLY_ENCLOSED_BY = '"' SKIP_HEADER = 1);"""
-    )
-
-    # Check if the data was loaded successfully
-    if cur.fetchall()[0][1] == "LOADED":
-        print({"copied_data_from_stage_into_table": True})
-    else:
+    try:
+        cur.execute(
+            f"""
+                        COPY INTO {table_name}
+                        FROM @{stage_name}
+                        FILE_FORMAT = (TYPE = CSV FIELD_OPTIONALLY_ENCLOSED_BY = '"' SKIP_HEADER = 1);"""
+        )
+
+        # Check if the data was loaded successfully
+        if cur.fetchall()[0][1] == "LOADED":
+            print({"copied_data_from_stage_into_table": True})
+        else:
+            print({"copied_data_from_stage_into_table": False})
+            sys.exit(1)
+    except Exception as e:
+        logger.error(f"Error copying data from stage to table: {e}")
         print({"copied_data_from_stage_into_table": False})
+        sys.exit(1)
 
 
 def put_file_to_stage(
@@ -119,25 +142,37 @@ def put_file_to_stage(
         stage_name (str): The name of the stage where the file will be uploaded.
         cur (snowflake.connector.cursor.SnowflakeCursor): The cursor to execute the SQL command.
     """
-    response = cur.execute(
-        f"PUT file://{file_name} @{stage_name} AUTO_COMPRESS=TRUE"
-    ).fetchall()
-    logger.error(response)
-
-    if response[0][6] == "UPLOADED":
-        print({"PUT_operation": True})
-    else:
+    try:
+        response = cur.execute(
+            f"PUT file://{file_name} @{stage_name} AUTO_COMPRESS=TRUE"
+        ).fetchall()
+        logger.error(response)
+
+        if response[0][6] == "UPLOADED":
+            print({"PUT_operation": True})
+        else:
+            print({"PUT_operation": False})
+            sys.exit(1)
+    except Exception as e:
+        logger.error(f"Error uploading file to stage: {e}")
         print({"PUT_operation": False})
+        sys.exit(1)
 
 
 def count_data_from_table(
     table_name: str, num_records: int, cur: snowflake.connector.cursor.SnowflakeCursor
 ):
-    count = cur.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0]
-    if count == num_records:
-        print({"data_transferred_completely": True})
-    else:
+    try:
+        count = cur.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0]
+        if count == num_records:
+            print({"data_transferred_completely": True})
+        else:
+            print({"data_transferred_completely": False})
+            sys.exit(1)
+    except Exception as e:
+        logger.error(f"Error counting data from table: {e}")
         print({"data_transferred_completely": False})
+        sys.exit(1)
 
 
 def compare_fetched_data(
@@ -157,21 +192,25 @@ def compare_fetched_data(
         repetitions (int): Number of times to repeat the comparison. Default is 10.
         fetch_limit (int): Number of rows to fetch from the table for comparison. Default is 100.
     """
-
-    fetched_data = cur.execute(
-        f"SELECT * FROM {table_name} LIMIT {fetch_limit}"
-    ).fetchall()
-
-    with open(file_name, newline="", encoding="utf-8") as csvfile:
-        reader = csv.reader(csvfile)
-        csv_data = list(reader)[1:]  # Skip header row
-        for _ in range(repetitions):
-            random_index = random.randint(0, fetch_limit - 1)
-            for y in range(len(fetched_data[0])):
-                if str(fetched_data[random_index][y]) != csv_data[random_index][y]:
-                    print({"data_integrity_check": False})
-                    break
-        print({"data_integrity_check": True})
+    try:
+        fetched_data = cur.execute(
+            f"SELECT * FROM {table_name} LIMIT {fetch_limit}"
+        ).fetchall()
+
+        with open(file_name, newline="", encoding="utf-8") as csvfile:
+            reader = csv.reader(csvfile)
+            csv_data = list(reader)[1:]  # Skip header row
+            for _ in range(repetitions):
+                random_index = random.randint(0, fetch_limit - 1)
+                for y in range(len(fetched_data[0])):
+                    if str(fetched_data[random_index][y]) != csv_data[random_index][y]:
+                        print({"data_integrity_check": False})
+                        break
+            print({"data_integrity_check": True})
+    except Exception as e:
+        logger.error(f"Error comparing fetched data: {e}")
+        print({"data_integrity_check": False})
+        sys.exit(1)
 
 
 def execute_get_command(stage_name: str, conn: snowflake.connector.SnowflakeConnection):
@@ -194,7 +233,12 @@ def execute_get_command(stage_name: str, conn: snowflake.connector.SnowflakeConn
             print({"GET_operation": True})
         else:
             print({"GET_operation": False})
+            sys.exit(1)
 
+    except Exception as e:
+        logger.error(f"Error downloading file from stage: {e}")
+        print({"GET_operation": False})
+        sys.exit(1)
     finally:
         try:
             for file in os.listdir(download_dir):
@@ -206,6 +250,7 @@ def execute_get_command(stage_name: str, conn: snowflake.connector.SnowflakeConn
             logger.error(
                 f"Error cleaning up directory {download_dir}. It may not exist or be empty."
             )
+            sys.exit(1)
 
 
 def perform_put_fetch_get(connection_parameters: dict, num_records: int = 1000):
@@ -252,20 +297,25 @@ def perform_put_fetch_get(connection_parameters: dict, num_records: int = 1000):
                 logger.error("Performing GET operation")
                 execute_get_command(stage_name, conn)
                 logger.error("File downloaded from stage to local directory")
-
     except Exception as e:
-        logger.error(f"Error during PUT/GET operation: {e}")
+        logger.error(f"Error during PUT_FETCH_GET operation: {e}")
+        sys.exit(1)
 
     finally:
+        logger.error("INFORMATIONING")
+        try:
         # Cleanup: Remove data from the stage and delete table
-        with connect(connection_parameters) as conn:
-            with conn.cursor() as cur:
-                cur.execute(f"REMOVE @{stage_name}")
-                cur.execute(f"DROP TABLE {table_name}")
+            with connect(connection_parameters) as conn:
+                with conn.cursor() as cur:
+                    cur.execute(f"REMOVE @{stage_name}")
+                    cur.execute(f"DROP TABLE {table_name}")
+            sys.exit(0)
+        except Exception as e:
+            logger.error(f"Error during cleanup: {e}")
+            sys.exit(1)
 
 
-# Disabled in MVP, uncomment to run
-# @prober_function
+@prober_function
 def perform_put_fetch_get_100_lines(connection_parameters: dict):
     """
     Performs a PUT and GET operation for 1,000 rows using the provided connection parameters.

prober/testing_matrix.json

@@ -1,16 +1,12 @@
 {
   "python-version": [
     {
-      "version": "3.10.17",
-      "snowflake-connector-python": ["3.14.0", "3.13.2", "3.13.1"]
+      "version": "3.13.4",
+      "snowflake-connector-python": ["3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]
     },
     {
       "version": "3.9.22",
-      "snowflake-connector-python": ["3.14.0", "3.13.2", "3.13.1"]
-    },
-    {
-      "version": "3.8.20",
-      "snowflake-connector-python": ["3.14.0", "3.13.2", "3.13.1"]
+      "snowflake-connector-python": ["3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]
     }
   ]
 }