Merge branch 'main' into SNOW-2043816-Kerberos-Proxy-Auth-Python

Author: sfc-gh-fpawlowski

DESCRIPTION.md

@@ -10,6 +10,8 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 - v3.16(TBD)
   - Bumped numpy dependency from <2.1.0 to <=2.2.4
   - Added Windows support for Python 3.13.
+  - Add `bulk_upload_chunks` parameter to `write_pandas` function. Setting this parameter to True changes the behaviour of write_pandas function to first write all the data chunks to the local disk and then perform the wildcard upload of the chunks folder to the stage. In default behaviour the chunks are being saved, uploaded and deleted one by one.
+
 
 - v3.15.1(May 20, 2025)
   - Added basic arrow support for Interval types.

prober/Dockerfile

@@ -0,0 +1,77 @@
+FROM alpine:3.18
+
+# boilerplate labels required by validation when pushing to ACR, ECR & GCR
+LABEL org.opencontainers.image.source="https://github.com/snowflakedb/snowflake-connector-python"
+LABEL com.snowflake.owners.email="[email protected]"
+LABEL com.snowflake.owners.slack="triage-snow-drivers-warsaw-dl"
+LABEL com.snowflake.owners.team="Snow Drivers"
+LABEL com.snowflake.owners.jira_area="Developer Platform"
+LABEL com.snowflake.owners.jira_component="Python Driver"
+# fake layers label to pass the validation
+LABEL com.snowflake.ugcbi.layers="sha256:850959b749c07b254308a4d1a84686fd7c09fcb94aeae33cc5748aa07e5cb232,sha256:b79d3c4628a989cbb8bc6f0bf0940ff33a68da2dca9c1ffbf8cfb2a27ac8d133,sha256:1cbcc0411a84fbce85e7ee2956c8c1e67b8e0edc81746a33d9da48c852037c3e,sha256:07e89b796f91d37255c6eec926b066d6818f3f2edc344a584d1b9566f77e1c27,sha256:84ff92691f909a05b224e1c56abb4864f01b4f8e3c854e4bb4c7baf1d3f6d652,sha256:3ab72684daee4eea64c3ae78a43ea332b86358446b6f2904dca4b634712e1537"
+
+RUN apk add --no-cache \
+    bash \
+    git \
+    make \
+    g++ \
+    zlib-dev \
+    openssl-dev \
+    libffi-dev \
+    jq
+
+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}
+
+# 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 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 from ARG MATRIX_VERSION
+RUN eval "$(pyenv init --path)" && \
+    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 $(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 setup.py prober
+COPY entrypoint.sh prober
+COPY probes/* prober/probes
+
+# Install /prober in editable mode for each virtual environment
+RUN for venv in ${HOME}/venvs/*; do \
+    source $venv/bin/activate && \
+    pip install -e ${HOME}/prober && \
+    deactivate; \
+done

prober/Jenkinsfile.groovy

@@ -0,0 +1,65 @@
+pipeline {
+    agent { label 'regular-memory-node' }
+
+    options {
+        ansiColor('xterm')
+        timestamps()
+    }
+
+    environment {
+        VAULT_CREDENTIALS = credentials('vault-jenkins')
+        COMMIT_SHA_SHORT = sh(script: 'cd PythonConnector/prober && git rev-parse --short HEAD', returnStdout: true).trim()
+        IMAGE_NAME = 'snowdrivers/python-driver-prober'
+        TEAM_NAME = 'Snow Drivers'
+        TEAM_JIRA_DL = 'triage-snow-drivers-warsaw-dl'
+        TEAM_JIRA_AREA = 'Developer Platform'
+        TEAM_JIRA_COMPONENT = 'Python Driver'
+    }
+
+    stages {
+        stage('Build Image') {
+            steps {
+                dir('./PythonConnector/prober') {
+                    sh """
+                    ls -l
+                    docker build \
+                    -t ${IMAGE_NAME}:${COMMIT_SHA_SHORT} \
+                    --label "org.opencontainers.image.revision=${COMMIT_SHA_SHORT}" \
+                    -f ./Dockerfile .
+                    """
+                }
+            }
+        }
+
+        stage('Checkout Jenkins Push Scripts') {
+            steps {
+                dir('k8sc-jenkins_scripts') {
+                    git branch: 'master',
+                    credentialsId: 'jenkins-snowflake-github-app-3',
+                    url: 'https://github.com/snowflakedb/k8sc-jenkins_scripts.git'
+                }
+            }
+        }
+
+        stage('Push Image') {
+            steps {
+                sh """
+                ./k8sc-jenkins_scripts/jenkins_push.sh \
+                -r "${VAULT_CREDENTIALS_USR}" \
+                -s "${VAULT_CREDENTIALS_PSW}" \
+                -i "${IMAGE_NAME}" \
+                -v "${COMMIT_SHA_SHORT}" \
+                -n "${TEAM_JIRA_DL}" \
+                -a "${TEAM_JIRA_AREA}" \
+                -C "${TEAM_JIRA_COMPONENT}"
+                """
+            }
+        }
+    }
+
+    post {
+        always {
+            cleanWs()
+        }
+    }
+}

prober/entrypoint.sh

@@ -0,0 +1,45 @@
+#!/bin/bash
+
+# Initialize an empty string to hold all parameters
+python_version=""
+connector_version=""
+params=""
+
+# Parse command-line arguments
+while [[ "$#" -gt 0 ]]; do
+    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
+
+# 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
 
@@ -29,12 +31,11 @@ def connect(connection_parameters: dict):
             database=connection_parameters["database"],
             schema=connection_parameters["schema"],
             role=connection_parameters["role"],
-            authenticator="KEY_PAIR_AUTHENTICATOR",
+            authenticator=connection_parameters["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}")
 
 
@@ -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.info(result)
         assert result == (1,)
-        logger.info({f"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:
-        logger.info({f"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,6 +1,8 @@
 import argparse
 import logging
+import sys
 
+from probes import login, put_fetch_get  # noqa
 from probes.logging_config import initialize_logger
 from probes.registry import PROBES_FUNCTIONS
 
@@ -19,30 +21,51 @@ 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="Database")
     parser.add_argument("--user", required=True, help="Username")
-    parser.add_argument("--private_key", required=True, help="Private key")
+    parser.add_argument(
+        "--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",
+    )
 
     # 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,
         "role": args.role,
         "account": args.account,
         "schema": args.schema,
         "warehouse": args.warehouse,
+        "database": args.database,
         "user": args.user,
-        "private_key": args.private_key,
+        "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.info("BBB")
-            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__":