snowpark-connect 0.27.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/snowpark_connect/utils/session.py

@@ -4,13 +4,17 @@
 
 import logging
 import os
+import threading
 from collections.abc import Sequence
 from typing import Any
 
 from snowflake import snowpark
-from snowflake.snowpark.exceptions import SnowparkClientException, SnowparkSQLException
+from snowflake.connector.description import PLATFORM
+from snowflake.snowpark.exceptions import SnowparkClientException
 from snowflake.snowpark.session import _get_active_session
 from snowflake.snowpark_connect.constants import DEFAULT_CONNECTION_NAME
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.utils.describe_query_cache import (
     instrument_session_for_describe_cache,
 )
@@ -21,6 +25,13 @@ from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import telemetry
 from snowflake.snowpark_connect.utils.udf_cache import init_builtin_udf_cache
 
+SKIP_SESSION_CONFIGURATION = False
+
+
+def skip_session_configuration(skip: bool):
+    global SKIP_SESSION_CONFIGURATION
+    SKIP_SESSION_CONFIGURATION = skip
+
 
 # Suppress experimental warnings from snowflake.snowpark logger
 def _filter_experimental_warnings(record):
@@ -50,7 +61,12 @@ def _get_current_snowpark_session() -> snowpark.Session | None:
 
 def configure_snowpark_session(session: snowpark.Session):
     """Configure a snowpark session with required parameters and settings."""
-    from snowflake.snowpark_connect.config import global_config
+    from snowflake.snowpark_connect.config import (
+        get_cte_optimization_enabled,
+        global_config,
+    )
+
+    global SKIP_SESSION_CONFIGURATION
 
     logger.info(f"Configuring session {session}")
 
@@ -63,11 +79,28 @@ def configure_snowpark_session(session: snowpark.Session):
     # custom udf imports
     session._python_files = set()
     session._import_files = set()
+    session._artifact_jars = set()
+
+    # custom artifact attributes
+    # track current chunk
+    # key: session_id, value: dict of (name, num_chunks, current_chunk_index)
+    session._current_chunk: dict[str, dict] = {}
+    # Use thread-safe access when modifying current chunk dictionary
+    session._current_chunk_lock = threading.RLock()
+
+    # track filenames to be uploaded to stage
+    # key: session_id, value: dict of (name, filename)
+    session._filenames: dict[str, dict[str, str]] = {}
+    # Use thread-safe access when modifying filenames dictionary
+    session._filenames_lock = threading.RLock()
 
     # built-in udf cache
     init_builtin_udf_cache(session)
     init_external_udxf_cache(session)
 
+    # file format cache
+    session._file_formats = set()
+
     # Set experimental parameters (warnings globally suppressed)
     session.ast_enabled = False
     session.eliminate_numeric_sql_value_cast_enabled = False
@@ -77,6 +110,18 @@ def configure_snowpark_session(session: snowpark.Session):
     session.connection.arrow_number_to_decimal_setter = True
     session.custom_package_usage_config["enabled"] = True
 
+    # Scoped temp objects may not be accessible in stored procedure and cause "object does not exist" error. So disable
+    # _use_scoped_temp_objects here and use temp table instead.
+    session._use_scoped_temp_objects = False
+
+    # Configure CTE optimization based on session configuration
+    cte_optimization_enabled = get_cte_optimization_enabled()
+    session.cte_optimization_enabled = cte_optimization_enabled
+    logger.info(f"CTE optimization enabled: {cte_optimization_enabled}")
+
+    # Default query tag to be used unless overridden by user using AppName or spark.addTag()
+    query_tag = "SNOWPARK_CONNECT_QUERY"
+
     default_fallback_timezone = "UTC"
     if global_config.spark_sql_session_timeZone is None:
         try:
@@ -103,35 +148,20 @@ def configure_snowpark_session(session: snowpark.Session):
         "TIMEZONE": f"'{global_config.spark_sql_session_timeZone}'",
         "QUOTED_IDENTIFIERS_IGNORE_CASE": "false",
         "PYTHON_SNOWPARK_ENABLE_THREAD_SAFE_SESSION": "true",
-        "PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS": "false",  # this is required for creating udfs from sproc
+        "ENABLE_STRUCTURED_TYPES_IN_SNOWPARK_CONNECT_RESPONSE": "true",
+        "QUERY_TAG": f"'{query_tag}'",
     }
 
-    session.sql(
-        f"ALTER SESSION SET {', '.join([f'{k} = {v}' for k, v in session_params.items()])}"
-    ).collect()
-
-    # Rolling ahead in preparation of GS release 9.22 (ETA 8/5/2025). Once 9.22 is past rollback risk, merge this
-    # parameter with other in the session_params dictionary above
-    try:
+    # SNOW-2245971: Stored procedures inside Native Apps run as Execute As Owner and hence cannot set session params.
+    if not SKIP_SESSION_CONFIGURATION:
         session.sql(
-            "ALTER SESSION SET ENABLE_STRUCTURED_TYPES_IN_SNOWPARK_CONNECT_RESPONSE=true"
+            f"ALTER SESSION SET {', '.join([f'{k} = {v}' for k, v in session_params.items()])}"
         ).collect()
-    except SnowparkSQLException:
-        logger.debug(
-            "ENABLE_STRUCTURED_TYPES_IN_SNOWPARK_CONNECT_RESPONSE is not defined"
+    else:
+        session_param_names = ", ".join(session_params.keys())
+        logger.info(
+            f"Skipping Snowpark Connect session configuration as requested. Please make sure following session parameters are set correctly: {session_param_names}"
         )
-    try:
-        session.sql(
-            "ALTER SESSION SET ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT=true"
-        ).collect()
-    except SnowparkSQLException:
-        logger.debug("ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT is not defined")
-    try:
-        session.sql(
-            "ALTER SESSION SET ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE=true"
-        ).collect()
-    except SnowparkSQLException:
-        logger.debug("ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE is not defined")
 
     # Instrument the snowpark session to use a cache for describe queries.
     instrument_session_for_describe_cache(session)
@@ -145,6 +175,10 @@ def _is_running_in_SPCS():
     )
 
 
+def _is_running_in_stored_procedure_or_notebook():
+    return PLATFORM == "XP"
+
+
 def _get_session_configs_from_ENV() -> dict[str, Any]:
     session_configs = {
         "account": os.getenv("SNOWFLAKE_ACCOUNT"),
@@ -198,11 +232,13 @@ def set_query_tags(spark_tags: Sequence[str]) -> None:
     """Sets Snowpark session query_tag value to the tag from the Spark request."""
 
     if any("," in tag for tag in spark_tags):
-        raise ValueError("Tags cannot contain ','.")
+        exception = ValueError("Tags cannot contain ','.")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
 
     # TODO: Tags might not be set correctly in parallel workloads or multi-threaded code.
     snowpark_session = get_or_create_snowpark_session()
     spark_tags_str = ",".join(sorted(spark_tags)) if spark_tags else None
 
-    if spark_tags_str != snowpark_session.query_tag:
+    if spark_tags_str and spark_tags_str != snowpark_session.query_tag:
         snowpark_session.query_tag = spark_tags_str

snowflake/snowpark_connect/utils/snowpark_connect_logging.py

@@ -6,17 +6,59 @@ import logging
 
 from pyspark import StorageLevel
 
-logger = logging.getLogger("snowflake_connect_server")
-logger.setLevel(logging.WARN)
 
-console_handler = logging.StreamHandler()
-console_handler.setLevel(logging.INFO)
-formatter = logging.Formatter(
-    "%(asctime)s - %(name)s - %(levelname)s - [Thread %(thread)d] - %(message)s"
+def ensure_logger_has_handler(
+    logger_name: str, log_level: int = logging.INFO, force_level: bool = False
+):
+    """
+    Ensure a logger has a StreamHandler, add one if missing.
+    Checks both the specific logger and root logger for existing handlers.
+
+    Args:
+        logger_name: Name of the logger to configure
+        log_level: Log level to set on both logger and handler
+        force_level: If True, always set the log level. If False, only set if logger level is NOTSET
+
+    Returns:
+        The configured logger
+    """
+    target_logger = logging.getLogger(logger_name)
+
+    # Only set level if forced or if logger hasn't been configured yet
+    if force_level or target_logger.level == logging.NOTSET:
+        target_logger.setLevel(log_level)
+    else:
+        log_level = target_logger.level
+
+    # Check if the logger already has a StreamHandler
+    has_stream_handler = any(
+        isinstance(h, logging.StreamHandler) for h in target_logger.handlers
+    )
+
+    # Check if root logger has handlers (from basicConfig or manual setup)
+    root_logger = logging.getLogger()
+    has_root_handlers = len(root_logger.handlers) > 0
+
+    # Only add handler if:
+    # 1. Logger doesn't have its own StreamHandler AND
+    # 2. Root logger doesn't have handlers (to avoid duplication)
+    if not has_stream_handler and not has_root_handlers:
+        handler = logging.StreamHandler()
+        handler.setLevel(log_level)
+        formatter = logging.Formatter(
+            "%(asctime)s - %(name)s - %(levelname)s - [Thread %(thread)d] - %(message)s"
+        )
+        handler.setFormatter(formatter)
+        target_logger.addHandler(handler)
+
+    return target_logger
+
+
+# Initialize the main logger using the helper function
+# force_level=False means it will respect any existing log level configuration
+logger = ensure_logger_has_handler(
+    "snowflake_connect_server", logging.INFO, force_level=False
 )
-console_handler.setFormatter(formatter)
-# Display the logs to the console
-logger.addHandler(console_handler)
 
 
 def run_once_decorator(func):

snowflake/snowpark_connect/utils/spcs_logger.py

@@ -0,0 +1,290 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+SPCS Logger - Adapted from ExecPlatform/src/coprocessor/python/telemetry/py/logger.py
+Outputs flat JSON format compatible with SPCS OpenTelemetry collector with proper trace context.
+"""
+
+import json
+import logging
+import sys
+import traceback
+from typing import Any, Mapping, Sequence
+
+
+class SPCSLoggerConfig:
+    """Configuration for SPCS logger."""
+
+    MESSAGE_SIZE_LIMIT_BYTES = 524288  # 512KB
+    ELLIPSIS = "..."
+
+    # Set to True if initialized
+    is_initialized = False
+
+
+def _encode_value_simple(value: Any) -> Any:
+    """
+    Encode a value to simple JSON format (not OpenTelemetry nested format).
+    SPCS expects flat JSON values, not the {stringValue: ...} format.
+    """
+    if isinstance(value, (bool, str, int, float)):
+        return value
+    if isinstance(value, Sequence) and not isinstance(value, str):
+        return [_encode_value_simple(v) for v in value]
+    if isinstance(value, Mapping):
+        return {str(k): _encode_value_simple(v) for k, v in value.items()}
+    # Stringify anything else
+    return str(value)
+
+
+# Skip Python's built-in LogRecord attributes
+_RESERVED_ATTRS = frozenset(
+    (
+        "asctime",
+        "args",
+        "created",
+        "exc_info",
+        "exc_text",
+        "filename",
+        "funcName",
+        "message",
+        "levelname",
+        "levelno",
+        "lineno",
+        "module",
+        "msecs",
+        "msg",
+        "name",
+        "pathname",
+        "process",
+        "processName",
+        "relativeCreated",
+        "stack_info",
+        "thread",
+        "threadName",
+        "taskName",
+    )
+)
+
+
+def _extract_attributes(record: logging.LogRecord) -> dict:
+    """Extract log record attributes to flat dict format for SPCS."""
+    attributes = {}
+
+    # Extract custom attributes from extra={}
+    for k, v in vars(record).items():
+        if k not in _RESERVED_ATTRS:
+            attributes[k] = _encode_value_simple(v)
+
+    # Add standard code location attributes
+    attributes["code.lineno"] = record.lineno
+    attributes["code.function"] = record.funcName
+    attributes["code.filepath"] = record.pathname
+
+    # Add exception info if present
+    if record.exc_info is not None:
+        exctype, value, tb = record.exc_info
+        if exctype is not None:
+            attributes["exception.type"] = exctype.__name__
+        if value is not None and value.args:
+            attributes["exception.message"] = str(value.args[0])
+        if tb is not None:
+            attributes["exception.stacktrace"] = "".join(
+                traceback.format_exception(*record.exc_info)
+            )
+
+    return attributes
+
+
+def get_snowflake_log_level_name(py_level_name: str) -> str:
+    """
+    Convert Python log level to Snowflake log level.
+    This matches the original UDF logger implementation.
+    """
+    level = py_level_name.upper()
+    if level == "WARNING":
+        return "WARN"
+    elif level == "CRITICAL":
+        return "FATAL"
+    elif level == "NOTSET":
+        return "TRACE"
+    else:
+        return level
+
+
+def get_severity_number(snowflake_level: str) -> int:
+    """
+    Get OTLP severity number (integer) for a Snowflake log level.
+
+    OTLP Spec: https://opentelemetry.io/docs/specs/otel/logs/data-model/#field-severitynumber
+    This returns INTEGER values (not strings like the buggy UDF code).
+    """
+    if snowflake_level == "TRACE":
+        return 1  # SEVERITY_NUMBER_TRACE
+    elif snowflake_level == "DEBUG":
+        return 5  # SEVERITY_NUMBER_DEBUG
+    elif snowflake_level == "INFO":
+        return 9  # SEVERITY_NUMBER_INFO
+    elif snowflake_level == "WARN":
+        return 13  # SEVERITY_NUMBER_WARN
+    elif snowflake_level == "ERROR":
+        return 17  # SEVERITY_NUMBER_ERROR
+    elif snowflake_level == "FATAL":
+        return 21  # SEVERITY_NUMBER_FATAL
+    else:
+        return 0  # SEVERITY_NUMBER_UNSPECIFIED
+
+
+def _encode_spcs_log_record(record: logging.LogRecord) -> dict:
+    """
+    Encode a log record to the FLAT JSON format expected by SPCS.
+
+    SPCS OpenTelemetry collector expects:
+    {
+      "body": "message",
+      "severity_text": "INFO",
+      "severity_number": 9,          # INTEGER, not string!
+      "attributes": {...},
+      "scope": {"name": "logger_name"}
+    }
+    """
+    # Format the message
+    message = str(record.msg)
+    if record.args:
+        try:
+            message = message % record.args
+        except (TypeError, ValueError):
+            message = str(record.msg)
+
+    # Truncate message if it exceeds size limit
+    message_bytes = message.encode("utf-8", errors="replace")
+    if sys.getsizeof(message_bytes) > SPCSLoggerConfig.MESSAGE_SIZE_LIMIT_BYTES:
+        truncate_length = SPCSLoggerConfig.MESSAGE_SIZE_LIMIT_BYTES - len(
+            SPCSLoggerConfig.ELLIPSIS.encode()
+        )
+        # Ensure we don't cut in the middle of a UTF-8 multibyte sequence
+        while truncate_length > 0 and (message_bytes[truncate_length] & 0xC0) == 0x80:
+            truncate_length -= 1
+        message_bytes = message_bytes[0:truncate_length]
+        message = (
+            message_bytes.decode("utf-8", errors="replace") + SPCSLoggerConfig.ELLIPSIS
+        )
+
+    # Map to Snowflake log level
+    snowflake_level = get_snowflake_log_level_name(record.levelname)
+
+    # Construct the FLAT log record (NOT nested OpenTelemetry structure)
+    log_record = {
+        "body": message,
+        "severity_text": snowflake_level,
+        "severity_number": get_severity_number(snowflake_level),  # INTEGER!
+        "attributes": _extract_attributes(record),
+        "scope": {"name": record.name},
+    }
+
+    return log_record
+
+
+# =============================================================================
+# SPCS-SPECIFIC HANDLER
+# =============================================================================
+
+
+class SPCSStreamHandler(logging.StreamHandler):
+    """
+    Custom handler for SPCS that writes flat JSON format to stdout.
+
+    The SPCS OpenTelemetry collector will:
+    1. Capture stdout
+    2. Parse JSON if line matches ^{.*}$
+    3. Extract body, severity_text, severity_number, attributes, scope, trace_id, span_id fields
+    4. Map trace_id/span_id to LogRecord protobuf fields
+    5. Backend creates TRACE column from protobuf trace_id/span_id
+    6. Route to Event Table
+    """
+
+    def __init__(self, stream=None) -> None:
+        """
+        Initialize the handler.
+
+        Args:
+            stream: Output stream (default: sys.stdout)
+        """
+        super().__init__(stream or sys.stdout)
+
+    def emit(self, record: logging.LogRecord):
+        """
+        Emit a log record as single-line JSON to stdout.
+        """
+        try:
+            # Encode to SPCS-compatible flat JSON format
+            log_record = _encode_spcs_log_record(record)
+
+            # Convert to compact JSON string (single line, no spaces)
+            log_json = json.dumps(log_record, separators=(",", ":"))
+
+            # Write to stdout (SPCS captures this)
+            self.stream.write(log_json + "\n")
+            self.flush()
+
+        except Exception:
+            self.handleError(record)
+
+
+# =============================================================================
+# INITIALIZATION FUNCTIONS
+# =============================================================================
+
+
+def setup_spcs_logger(
+    log_level: int = logging.INFO,
+    logger_name: str = None,
+    enable_console_output: bool = False,
+) -> logging.Logger:
+    """
+    Set up the root logger for SPCS with flat JSON formatting.
+
+    Args:
+        log_level: Python logging level (e.g., logging.INFO)
+        logger_name: Optional logger name (None for root logger)
+        enable_console_output: If True, also adds a human-readable console handler to stderr
+
+    Returns:
+        Configured logger instance
+
+    Example:
+        >>> logger = setup_spcs_logger(logging.INFO, enable_console_output=True)
+        >>> logger.info("Hello from SPCS", extra={"user_id": 123, "action": "login"})
+
+        # Output to stdout (captured by SPCS):
+        {"body":"Hello from SPCS","severity_text":"INFO","severity_number":9,"attributes":{"user_id":123,"action":"login","code.lineno":42,"code.function":"main","code.filepath":"/app/main.py"},"scope":{"name":"root"}}
+
+        # Output to stderr (if enable_console_output=True):
+        2024-01-15 10:30:45,123 - root - INFO - Hello from SPCS
+    """
+    # Mark as initialized
+    SPCSLoggerConfig.is_initialized = True
+
+    # Get logger (root or named)
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(log_level)
+    logger.handlers.clear()
+
+    # Add SPCS flat JSON handler (writes JSON to stdout)
+    spcs_handler = SPCSStreamHandler(sys.stdout)
+    spcs_handler.setLevel(log_level)
+    logger.addHandler(spcs_handler)
+
+    # Optionally add human-readable console handler (to stderr to avoid mixing with JSON logs)
+    if enable_console_output:
+        console_handler = logging.StreamHandler(sys.stderr)
+        console_handler.setLevel(log_level)
+        formatter = logging.Formatter(
+            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        )
+        console_handler.setFormatter(formatter)
+        logger.addHandler(console_handler)
+
+    return logger