snowpark-connect 0.21.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

snowflake/snowpark_connect/server.py

@@ -83,6 +83,7 @@ from snowflake.snowpark_connect.utils.context import (
     set_session_id,
     set_spark_version,
 )
+from snowflake.snowpark_connect.utils.env_utils import get_int_from_env
 from snowflake.snowpark_connect.utils.interrupt import (
     interrupt_all_queries,
     interrupt_queries_with_tag,
@@ -700,11 +701,27 @@ def _serve(
             return
 
         server_options = [
-            ("grpc.max_receive_message_length", _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE),
-            ("grpc.max_metadata_size", _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE),
+            (
+                "grpc.max_receive_message_length",
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_MESSAGE_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE,
+                ),
+            ),
+            (
+                "grpc.max_metadata_size",
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
+                ),
+            ),
             (
                 "grpc.absolute_max_metadata_size",
-                _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE * 2,
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
+                )
+                * 2,
             ),
         ]
         server = grpc.server(
@@ -812,8 +829,11 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
     Spark Connect gRPC channel builder for Unix domain sockets
     """
 
-    def __init__(self, channelOptions: Optional[List[Tuple[str, Any]]] = None) -> None:
-        url: str = get_client_url()
+    def __init__(
+        self, url: str = None, channelOptions: Optional[List[Tuple[str, Any]]] = None
+    ) -> None:
+        if url is None:
+            url = get_client_url()
         if url[:6] != "unix:/" or len(url) < 7:
             raise PySparkValueError(
                 error_class="INVALID_CONNECT_URL",
@@ -981,6 +1001,7 @@ def start_session(
     stop_event: threading.Event = None,
     snowpark_session: Optional[snowpark.Session] = None,
     connection_parameters: Optional[Dict[str, str]] = None,
+    max_grpc_message_size: int = _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE,
 ) -> threading.Thread | None:
     """
     Starts Spark Connect server connected to Snowflake. No-op if the Server is already running.
@@ -1003,6 +1024,14 @@ def start_session(
                                 provided, the `snowpark_session` parameter must be None.
     """
     try:
+        # Changing the value of our global variable based on the grpc message size provided by the user.
+        global _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE
+        _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE = max_grpc_message_size
+
+        from pyspark.sql.connect.client import ChannelBuilder
+
+        ChannelBuilder.MAX_MESSAGE_LENGTH = max_grpc_message_size
+
         if os.environ.get("SPARK_ENV_LOADED"):
             raise RuntimeError(
                 "Snowpark Connect cannot be run inside of a Spark environment"

snowflake/snowpark_connect/type_mapping.py

@@ -59,7 +59,7 @@ def _get_struct_type_class():
 
 
 @cache
-def _get_python_sql_utils_class():
+def get_python_sql_utils_class():
     return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
@@ -70,7 +70,7 @@ def parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
     This mimics pysparks.ddl parsing logic pyspark.sql.types._py_parse_datatype_string
     """
     struct_type_class = _get_struct_type_class()
-    python_sql_utils = _get_python_sql_utils_class()
+    python_sql_utils = get_python_sql_utils_class()
 
     try:
         # DDL format, "fieldname datatype, fieldname datatype".
@@ -324,6 +324,8 @@ def cast_to_match_snowpark_type(
             return str(content)
         case snowpark.types.VariantType:
             return str(content)
+        case snowpark.types.TimestampType:
+            return str(content)
         case _:
             raise SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type in casting: {data_type}"
@@ -779,6 +781,8 @@ def map_simple_types(simple_type: str) -> snowpark.types.DataType:
             return snowpark.types.TimestampType()
         case "timestamp_ntz":
             return snowpark.types.TimestampType(snowpark.types.TimestampTimeZone.NTZ)
+        case "timestamp_ltz":
+            return snowpark.types.TimestampType(snowpark.types.TimestampTimeZone.LTZ)
         case "day_time_interval":
             # this is not a column type in snowflake so there won't be a dataframe column
             # with this, for now this type won't make any sense

snowflake/snowpark_connect/utils/describe_query_cache.py

@@ -131,21 +131,14 @@ def instrument_session_for_describe_cache(session: snowpark.Session):
             logger.debug(f"DDL detected, clearing describe query cache: '{query}'")
             cache.clear()
 
-    def report_query(qid: str, is_internal: bool) -> None:
-        if is_internal:
-            telemetry.report_internal_query()
-        elif qid:
-            telemetry.report_query_id(qid)
-
     def wrap_execute(wrapped_fn):
         def fn(query: str, **kwargs):
             update_cache_for_query(query)
-            is_internal = kwargs.get("_is_internal", False)
             try:
                 result = wrapped_fn(query, **kwargs)
-                report_query(result.sfqid, is_internal)
+                telemetry.report_query(result, **kwargs)
             except Exception as e:
-                report_query(e.sfqid, is_internal)
+                telemetry.report_query(e, **kwargs)
                 raise e
             return result
 

snowflake/snowpark_connect/utils/env_utils.py

@@ -0,0 +1,55 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Environment variable utilities for Snowpark Connect.
+"""
+
+import os
+
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+
+
+def get_int_from_env(env_var: str, default: int) -> int:
+    """
+    Safely get integer value from environment variable with fallback to default.
+
+    Args:
+        env_var: Environment variable name
+        default: Default integer value if env var is not set or invalid
+
+    Returns:
+        Integer value from environment variable or default
+
+    Raises:
+        TypeError: If default is not an integer
+
+    Examples:
+        >>> get_int_from_env("MAX_WORKERS", 10)
+        10
+        >>> os.environ["MAX_WORKERS"] = "20"
+        >>> get_int_from_env("MAX_WORKERS", 10)
+        20
+        >>> os.environ["MAX_WORKERS"] = "invalid"
+        >>> get_int_from_env("MAX_WORKERS", 10)  # logs warning, returns 10
+        10
+    """
+    # Validate that default is actually an integer
+    if not isinstance(default, int):
+        raise TypeError(
+            f"Default value must be an integer, got {type(default).__name__}: {default}"
+        )
+
+    value = os.getenv(env_var)
+    if value is None:
+        return default
+
+    try:
+        return int(value)
+    except ValueError:
+        logger.warning(
+            f"Invalid integer value for environment variable {env_var}: '{value}', "
+            f"using default: {default}"
+        )
+        return default

snowflake/snowpark_connect/utils/session.py

@@ -73,6 +73,27 @@ def configure_snowpark_session(session: snowpark.Session):
     session.connection.arrow_number_to_decimal_setter = True
     session.custom_package_usage_config["enabled"] = True
 
+    default_fallback_timezone = "UTC"
+    if global_config.spark_sql_session_timeZone is None:
+        try:
+            result = session.sql("SHOW PARAMETERS LIKE 'TIMEZONE'").collect()
+            if result and len(result) > 0:
+                value = result[0]["value"]
+                logger.warning(
+                    f"Using Snowflake session timezone parameter as fallback: {value}"
+                )
+            else:
+                value = default_fallback_timezone
+                logger.warning(
+                    f"Could not determine timezone from parameters, defaulting to {default_fallback_timezone}"
+                )
+        except Exception as e:
+            value = default_fallback_timezone
+            logger.warning(
+                f"Could not query Snowflake timezone parameter ({e}), defaulting to {default_fallback_timezone}"
+            )
+        global_config.spark_sql_session_timeZone = value
+
     session_params = {
         "TIMESTAMP_TYPE_MAPPING": "TIMESTAMP_LTZ",
         "TIMEZONE": f"'{global_config.spark_sql_session_timeZone}'",
@@ -181,7 +202,3 @@ def set_query_tags(spark_tags: Sequence[str]) -> None:
 
     if spark_tags_str != snowpark_session.query_tag:
         snowpark_session.query_tag = spark_tags_str
-
-
-def get_python_udxf_import_files(session: snowpark.Session) -> str:
-    return ",".join([file for file in [*session._python_files, *session._import_files]])