SNOW-2114093 Probing script (#2335)

Author: sfc-gh-pcyrek

.gitignore

@@ -125,3 +125,7 @@ core.*
 # Compiled Cython
 src/snowflake/connector/arrow_iterator.cpp
 src/snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_arrow_iterator.cpp
+
+# Prober files
+prober/parameters.json
+prober/snowflake_prober.egg-info/

MANIFEST.in

@@ -20,7 +20,6 @@ exclude tox.ini
 exclude mypy.ini
 exclude .clang-format
 exclude .wiremock/*
-exclude cmd/prober/testing_matrix.json
 
 prune ci
 prune benchmark
@@ -29,3 +28,4 @@ prune tested_requirements
 prune src/snowflake/connector/nanoarrow_cpp/scripts
 prune __pycache__
 prune samples
+prune prober

prober/__init__.py

@@ -0,0 +1,30 @@
+import logging
+
+
+def initialize_logger(name=__name__, level=logging.INFO):
+    """
+    Initializes and configures a logger.
+
+    Args:
+        name (str): The name of the logger.
+        level (int): The logging level (e.g., logging.INFO, logging.DEBUG).
+
+    Returns:
+        logging.Logger: Configured logger instance.
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+
+    # Create a console handler
+    handler = logging.StreamHandler()
+    handler.setLevel(level)
+
+    # Create a formatter and set it for the handler
+    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+    handler.setFormatter(formatter)
+
+    # Add the handler to the logger
+    if not logger.handlers:  # Avoid duplicate handlers
+        logger.addHandler(handler)
+
+    return logger

prober/probes/__init__.py

@@ -0,0 +1,66 @@
+from probes.logging_config import initialize_logger
+from probes.registry import prober_function
+
+import snowflake.connector
+
+# Initialize logger
+logger = initialize_logger(__name__)
+
+
+def connect(connection_parameters: dict):
+    """
+    Initializes the Python driver for login using the provided connection parameters.
+
+    Args:
+        connection_parameters (dict): A dictionary containing connection details such as
+                                       host, port, user, password, account, schema, etc.
+
+    Returns:
+        snowflake.connector.SnowflakeConnection: A connection object if successful.
+    """
+    try:
+        # Initialize the Snowflake connection
+        connection = snowflake.connector.connect(
+            user=connection_parameters["user"],
+            account=connection_parameters["account"],
+            host=connection_parameters["host"],
+            port=connection_parameters["port"],
+            warehouse=connection_parameters["warehouse"],
+            database=connection_parameters["database"],
+            schema=connection_parameters["schema"],
+            role=connection_parameters["role"],
+            authenticator="KEY_PAIR_AUTHENTICATOR",
+            private_key=connection_parameters["private_key"],
+        )
+        return connection
+    except Exception as e:
+        logger.info({f"success_login={False}"})
+        logger.error(f"Error connecting to Snowflake: {e}")
+
+
+@prober_function
+def perform_login(connection_parameters: dict):
+    """
+    Performs the login operation using the provided connection parameters.
+
+    Args:
+        connection_parameters (dict): A dictionary containing connection details such as
+                                       host, port, user, password, account, schema, etc.
+
+    Returns:
+        bool: True if login is successful, False otherwise.
+    """
+    try:
+        # Connect to Snowflake
+        connection = connect(connection_parameters)
+
+        # Perform a simple query to test the connection
+        cursor = connection.cursor()
+        cursor.execute("SELECT 1;")
+        result = cursor.fetchone()
+        logger.info(result)
+        assert result == (1,)
+        logger.info({f"success_login={True}"})
+    except Exception as e:
+        logger.info({f"success_login={False}"})
+        logger.error(f"Error during login: {e}")

prober/probes/logging_config.py

@@ -0,0 +1,49 @@
+import argparse
+import logging
+
+from probes.logging_config import initialize_logger
+from probes.registry import PROBES_FUNCTIONS
+
+# Initialize logger
+logger = initialize_logger(__name__)
+
+
+def main():
+    logger.info("Starting Python Driver Prober...")
+    # Set up argument parser
+    parser = argparse.ArgumentParser(description="Python Driver Prober")
+    parser.add_argument("--scope", required=True, help="Scope of probing")
+    parser.add_argument("--host", required=True, help="Host")
+    parser.add_argument("--port", type=int, required=True, help="Port")
+    parser.add_argument("--role", required=True, help="Protocol")
+    parser.add_argument("--account", required=True, help="Account")
+    parser.add_argument("--schema", required=True, help="Schema")
+    parser.add_argument("--warehouse", required=True, help="Warehouse")
+    parser.add_argument("--user", required=True, help="Username")
+    parser.add_argument("--private_key", required=True, help="Private key")
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    connection_params = {
+        "host": args.host,
+        "port": args.port,
+        "role": args.role,
+        "account": args.account,
+        "schema": args.schema,
+        "warehouse": args.warehouse,
+        "user": args.user,
+        "private_key": args.private_key,
+    }
+
+    for function_name, function in PROBES_FUNCTIONS.items():
+        try:
+            logging.info("BBB")
+            logging.error(f"Running probe: {function_name}")
+            function(connection_params)
+        except Exception as e:
+            logging.error(f"Error running probe {function_name}: {e}")
+
+
+if __name__ == "__main__":
+    main()

prober/probes/login.py

@@ -0,0 +1,10 @@
+PROBES_FUNCTIONS = {}
+
+
+def prober_function(func):
+    """
+    Register a function in the PROBES_FUNCTIONS dictionary.
+    The key is the function name, and the value is the function itself.
+    """
+    PROBES_FUNCTIONS[func.__name__] = func
+    return func

prober/probes/main.py

@@ -0,0 +1,16 @@
+from setuptools import find_packages, setup
+
+setup(
+    name="snowflake_prober",
+    version="1.0.0",
+    packages=find_packages(),
+    install_requires=[
+        "snowflake-connector-python",
+        "requests",
+    ],
+    entry_points={
+        "console_scripts": [
+            "prober=probes.main:main",
+        ],
+    },
+)