[SNOW-2114104] Config generator and update dockerfile (#2350)

Co-authored-by: Maxim Mishchenko <[email protected]>

Author: sfc-gh-pcyrek

prober/Dockerfile

@@ -17,48 +17,61 @@ RUN apk add --no-cache \
     g++ \
     zlib-dev \
     openssl-dev \
-    libffi-dev
+    libffi-dev \
+    jq
 
-ENV HOME="/root"
+ENV HOME="/home/driveruser"
+
+# Create a group with GID=1000 and a user with UID=1000
+RUN addgroup -g 1000 drivergroup && \
+    adduser -u 1000 -G drivergroup -D driveruser
+
+# Set permissions for the non-root user
+RUN mkdir -p ${HOME} && \
+    chown -R driveruser:drivergroup ${HOME}
+
+# Switch to the non-root user
+USER driveruser
 WORKDIR ${HOME}
-RUN git clone --depth=1 https://github.com/pyenv/pyenv.git .pyenv
+
+# Set environment variables
 ENV PYENV_ROOT="${HOME}/.pyenv"
 ENV PATH="${PYENV_ROOT}/shims:${PYENV_ROOT}/bin:${PATH}"
 
 
+# Install pyenv
+RUN git clone --depth=1 https://github.com/pyenv/pyenv.git ${PYENV_ROOT}
+
 # Build arguments for Python versions and Snowflake connector versions
-ARG PYTHON_VERSIONS="3.8.20 3.9.22 3.10.17"
-ARG SNOWFLAKE_CONNECTOR_VERSIONS="3.14.0 3.13.2 3.13.1"
+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"]}'
 
 
-# Install Python versions
+# Install Python versions from ARG MATRIX_VERSION
 RUN eval "$(pyenv init --path)" && \
-    for version in $PYTHON_VERSIONS; do \
-        pyenv install $version || echo "Failed to install Python $version"; \
+    for python_version in $(echo $MATRIX_VERSION | jq -r 'keys[]'); do \
+        pyenv install $python_version || echo "Failed to install Python $python_version"; \
     done
 
-
 # Create virtual environments for each combination of Python and Snowflake connector versions
-RUN for python_version in $PYTHON_VERSIONS; do \
-    for connector_version in $SNOWFLAKE_CONNECTOR_VERSIONS; do \
-        venv_path="/venvs/python_${python_version}_connector_${connector_version}"; \
+RUN for python_version in $(echo $MATRIX_VERSION | jq -r 'keys[]'); do \
+    for connector_version in $(echo $MATRIX_VERSION | jq -r ".\"${python_version}\"[]"); do \
+        venv_path="${HOME}/venvs/python_${python_version}_connector_${connector_version}"; \
         $PYENV_ROOT/versions/$python_version/bin/python -m venv $venv_path && \
         $venv_path/bin/pip install --upgrade pip && \
         $venv_path/bin/pip install snowflake-connector-python==$connector_version; \
     done; \
 done
 
 # Copy the prober script into the container
-RUN mkdir -p /prober/probes/
-COPY __init__.py /prober
-# COPY parameters.json /prober
-COPY setup.py /prober
-COPY entrypoint.sh /prober
-COPY probes/* /prober/probes
+RUN mkdir -p prober/probes/
+COPY __init__.py prober
+COPY setup.py prober
+COPY entrypoint.sh prober
+COPY probes/* prober/probes
 
 # Install /prober in editable mode for each virtual environment
-RUN for venv in /venvs/*; do \
+RUN for venv in ${HOME}/venvs/*; do \
     source $venv/bin/activate && \
-    pip install -e /prober && \
+    pip install -e ${HOME}/prober && \
     deactivate; \
 done

prober/entrypoint.sh

@@ -1,18 +1,45 @@
 #!/bin/bash
 
 # Initialize an empty string to hold all parameters
+python_version=""
+connector_version=""
 params=""
 
-# Parse command-line arguments dynamically
+# Parse command-line arguments
 while [[ "$#" -gt 0 ]]; do
-    params="$params $1 $2"
-    shift 2
+    if [[ "$1" == "--python_version" ]]; then
+        python_version="$2"
+        shift 2
+    elif [[ "$1" == "--connector_version" ]]; then
+        connector_version="$2"
+        shift 2
+    else
+        params+="$1 $2 "
+        shift 2
+    fi
 done
 
-# Run main.py with all available virtual environments
-for venv in /venvs/*; do
-    echo "Running main.py with virtual environment: $(basename "$venv")"
-    source "$venv/bin/activate"
-    prober $params
-    deactivate
-done
+# Construct the virtual environment path
+venv_path="${HOME}/venvs/python_${python_version}_connector_${connector_version}"
+
+# Check if the virtual environment exists
+if [[ ! -d "$venv_path" ]]; then
+    echo "Error: Virtual environment not found at $venv_path"
+    exit 1
+fi
+
+# Run main.py with given venv
+echo "Running main.py with virtual environment: $venv_path"
+source "$venv_path/bin/activate"
+prober $params
+status=$?
+deactivate
+
+# Check the exit status of prober
+if [[ $status -ne 0 ]]; then
+    echo "Error: prober returned failure."
+    exit 1
+else
+    echo "Success: prober returned success."
+    exit 0
+fi

prober/probes/login.py

@@ -1,3 +1,5 @@
+import sys
+
 from probes.logging_config import initialize_logger
 from probes.registry import prober_function
 
@@ -30,11 +32,10 @@ def connect(connection_parameters: dict):
             schema=connection_parameters["schema"],
             role=connection_parameters["role"],
             authenticator=connection_parameters["authenticator"],
-            private_key_file=connection_parameters["private_key_file"],
+            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}")
 
 
@@ -54,13 +55,22 @@ def perform_login(connection_parameters: dict):
         # Connect to Snowflake
         connection = connect(connection_parameters)
 
+        # Log the connection details
+        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+        driver_version = snowflake.connector.__version__
+
         # Perform a simple query to test the connection
         cursor = connection.cursor()
         cursor.execute("SELECT 1;")
         result = cursor.fetchone()
-        logger.error(f"Logging: {result}")
         assert result == (1,)
-        print({"success_login": True})
+        print(
+            f"cloudprober_driver_python_perform_login{{python_version={python_version}, driver_version={driver_version}}} 0"
+        )
+        sys.exit(0)
     except Exception as e:
-        print({"success_login": False})
+        print(
+            f"cloudprober_driver_python_perform_login{{python_version={python_version}, driver_version={driver_version}}} 1"
+        )
         logger.error(f"Error during login: {e}")
+        sys.exit(1)

prober/probes/main.py

@@ -1,5 +1,6 @@
 import argparse
 import logging
+import sys
 
 from probes import login, put_fetch_get  # noqa
 from probes.logging_config import initialize_logger
@@ -23,13 +24,23 @@ def main():
     parser.add_argument("--database", required=True, help="Database")
     parser.add_argument("--user", required=True, help="Username")
     parser.add_argument(
-        "--auth", required=True, help="Authenticator (e.g., KEY_PAIR_AUTHENTICATOR)"
+        "--authenticator",
+        required=True,
+        help="Authenticator (e.g., KEY_PAIR_AUTHENTICATOR)",
+    )
+    parser.add_argument(
+        "--private_key_file",
+        required=True,
+        help="Private key file in DER format base64-encoded and '/' -> '_', '+' -> '-' replacements",
     )
-    parser.add_argument("--private_key_file", required=True, help="Private key pwd")
 
     # Parse arguments
     args = parser.parse_args()
 
+    private_key = (
+        open(args.private_key_file).read().strip().replace("_", "/").replace("-", "+")
+    )
+
     connection_params = {
         "host": args.host,
         "port": args.port,
@@ -39,16 +50,22 @@ def main():
         "warehouse": args.warehouse,
         "database": args.database,
         "user": args.user,
-        "authenticator": args.auth,
-        "private_key_file": args.private_key_file,
+        "authenticator": args.authenticator,
+        "private_key": private_key,
     }
 
-    for function_name, function in PROBES_FUNCTIONS.items():
+    if args.scope not in PROBES_FUNCTIONS:
+        logging.error(
+            f"Invalid scope: {args.scope}. Available scopes: {list(PROBES_FUNCTIONS.keys())}"
+        )
+        sys.exit(1)
+    else:
+        logging.info(f"Running probe for scope: {args.scope}")
         try:
-            logging.error(f"Running probe: {function_name}")
-            function(connection_params)
+            PROBES_FUNCTIONS[args.scope](connection_params)
         except Exception as e:
-            logging.error(f"Error running probe {function_name}: {e}")
+            logging.error(f"Error running probe {args.scope}: {e}")
+            sys.exit(1)
 
 
 if __name__ == "__main__":

prober/probes/testing_matrix.json

@@ -0,0 +1,16 @@
+{
+  "python-version": [
+    {
+      "version": "3.10.17",
+      "snowflake-connector-python": ["3.14.0", "3.13.2", "3.13.1"]
+    },
+    {
+      "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"]
+    }
+  ]
+}

prober/testing_matrix.json

@@ -0,0 +1,31 @@
+import json
+
+
+def extract_versions():
+    with open("testing_matrix.json") as file:
+        data = json.load(file)
+    version_mapping = {}
+    for entry in data["python-version"]:
+        python_version = str(entry["version"])
+        version_mapping[python_version] = entry["snowflake-connector-python"]
+    return version_mapping
+
+
+def update_dockerfile(version_mapping):
+    dockerfile_path = "Dockerfile"
+    new_matrix_version = json.dumps(version_mapping)
+
+    with open(dockerfile_path) as file:
+        lines = file.readlines()
+
+    with open(dockerfile_path, "w") as file:
+        for line in lines:
+            if line.startswith("ARG MATRIX_VERSION"):
+                file.write(f"ARG MATRIX_VERSION='{new_matrix_version}'\n")
+            else:
+                file.write(line)
+
+
+if __name__ == "__main__":
+    extracted_mapping = extract_versions()
+    update_dockerfile(extracted_mapping)