SNOW-2114096: Implementing prober script image builder (#2340)

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

Author: sfc-gh-pcyrek

prober/Dockerfile

@@ -0,0 +1,54 @@
+FROM alpine:3.18
+
+RUN apk add --no-cache \
+    bash \
+    git \
+    make \
+    g++ \
+    zlib-dev \
+    openssl-dev \
+    libffi-dev
+
+ENV HOME="/root"
+WORKDIR ${HOME}
+RUN git clone --depth=1 https://github.com/pyenv/pyenv.git .pyenv
+ENV PYENV_ROOT="${HOME}/.pyenv"
+ENV PATH="${PYENV_ROOT}/shims:${PYENV_ROOT}/bin:${PATH}"
+
+
+# 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"
+
+
+# Install Python versions
+RUN eval "$(pyenv init --path)" && \
+    for version in $PYTHON_VERSIONS; do \
+        pyenv install $version || echo "Failed to install 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}"; \
+        $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
+
+# Install /prober in editable mode for each virtual environment
+RUN for venv in /venvs/*; do \
+    source $venv/bin/activate && \
+    pip install -e /prober && \
+    deactivate; \
+done

prober/entrypoint.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+# Initialize an empty string to hold all parameters
+params=""
+
+# Parse command-line arguments dynamically
+while [[ "$#" -gt 0 ]]; do
+    params="$params $1 $2"
+    shift 2
+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

prober/probes/login.py

@@ -29,8 +29,8 @@ def connect(connection_parameters: dict):
             database=connection_parameters["database"],
             schema=connection_parameters["schema"],
             role=connection_parameters["role"],
-            authenticator="KEY_PAIR_AUTHENTICATOR",
-            private_key=connection_parameters["private_key"],
+            authenticator=connection_parameters["authenticator"],
+            private_key_file=connection_parameters["private_key_file"],
         )
         return connection
     except Exception as e:

prober/probes/main.py

@@ -1,6 +1,7 @@
 import argparse
 import logging
 
+from probes import login  # noqa
 from probes.logging_config import initialize_logger
 from probes.registry import PROBES_FUNCTIONS
 
@@ -19,8 +20,12 @@ def main():
     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("--database", required=True, help="Datanase")
     parser.add_argument("--user", required=True, help="Username")
-    parser.add_argument("--private_key", required=True, help="Private key")
+    parser.add_argument(
+        "--auth", required=True, help="Authenticator (e.g., KEY_PAIR_AUTHENTICATOR)"
+    )
+    parser.add_argument("--private_key_file", required=True, help="Private key pwd")
 
     # Parse arguments
     args = parser.parse_args()
@@ -32,13 +37,14 @@ def main():
         "account": args.account,
         "schema": args.schema,
         "warehouse": args.warehouse,
+        "database": args.database,
         "user": args.user,
-        "private_key": args.private_key,
+        "authenticator": args.auth,
+        "private_key_file": args.private_key_file,
     }
 
     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: