snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/config.py

@@ -0,0 +1,576 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+# Proto source for reference:
+# https://github.com/apache/spark/blob/branch-3.5/connector/connect/common/src/main/protobuf/spark/connect/base.proto#L420
+import os
+import re
+import sys
+import time
+from collections import defaultdict
+from copy import copy
+from typing import Any
+
+import jpype
+import pyspark.sql.connect.proto.base_pb2 as proto_base
+from tzlocal import get_localzone_name
+
+from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+)
+from snowflake.snowpark.exceptions import SnowparkSQLException
+from snowflake.snowpark_connect.utils.concurrent import SynchronizedDict
+from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+    telemetry,
+)
+from snowflake.snowpark_connect.version import VERSION as sas_version
+
+
+def str_to_bool(boolean_str: str) -> bool:
+    assert boolean_str in [
+        "True",
+        "true",
+        "False",
+        "false",
+        "1",
+        "0",
+        "",  # This is the default value, equivalent to False.
+    ], f"Invalid boolean value: {boolean_str}"
+    return boolean_str in ["True", "true", "1"]
+
+
+class GlobalConfig:
+    """This class contains the global configuration for the Spark Server."""
+
+    default_static_global_config = {
+        # Defaults from Spark https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+        "spark.sql.warehouse.dir": None,
+        "spark.sql.catalogImplementation": "in-memory",
+        "spark.sql.sources.schemaStringLengthThreshold": 4000,
+        "spark.sql.filesourceTableRelationCacheSize": 1000,
+        "spark.sql.codegen.cache.maxEntries": 100,
+        "spark.sql.codegen.comments": False,
+        "spark.sql.debug": False,
+        "spark.sql.hive.thriftServer.singleSession": False,
+        "spark.sql.extensions": None,
+        "spark.sql.cache.serializer": "org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer",
+        "spark.sql.queryExecutionListeners": 1000,
+        "spark.sql.shuffleExchange.maxThreadThreshold": 1024,
+        "spark.sql.broadcastExchange.maxThreadThreshold": 128,
+        "spark.sql.subquery.maxThreadThreshold": 16,
+        "spark.sql.resultQueryStage.maxThreadThreshold": 1024,
+        "spark.sql.event.truncate.length": sys.maxsize,
+        "spark.sql.legacy.sessionInitWithConfigDefaults": False,
+        "spark.sql.defaultUrlStreamHandlerFactory.enabled": True,
+        "spark.sql.streaming.ui.enabled": True,
+        "spark.sql.streaming.ui.retainedProgressUpdates": 100,
+        "spark.sql.streaming.ui.retainedQueries": 100,
+        "spark.sql.metadataCacheTTLSeconds": -1,
+        "spark.sql.streaming.ui.enabledCustomMetricList": None,
+        "spark.sql.sources.disabledJdbcConnProviderList": "",
+        "spark.python.sql.dataFrameDebugging.enabled": True,
+        "spark.sql.extensions.test.loadFromCp": None,
+        "spark.sql.streaming.streamingQueryListeners": None,
+        # Defaults from Spark https://github.com/apache/spark/blob/master/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
+        "spark.connect.grpc.binding.address": 15002,
+        "spark.connect.grpc.port.maxRetries": 0,
+        "spark.connect.grpc.interceptor.classes": 128 * 1024 * 1024,
+        "spark.connect.grpc.maxInboundMessageSize": 128 * 1024 * 1024,
+        "spark.connect.grpc.marshallerRecursionLimit": 1024,
+        "spark.connect.session.manager.defaultSessionTimeout": True,
+        "spark.connect.execute.reattachable.senderMaxStreamDuration": None,
+        "spark.connect.extensions.expression.classes": None,
+        "spark.connect.extensions.command.classes": None,
+        "spark.connect.jvmStacktrace.maxSize": 1024,
+        "spark.sql.connect.ui.retainedStatements": 200,
+        "spark.connect.copyFromLocalToFs.allowDestLocal": False,
+        "spark.sql.connect.ui.retainedSessions": 200,
+        "spark.connect.grpc.maxMetadataSize": 1024,
+        "spark.connect.session.planCache.maxSize": 16,
+        "spark.connect.authenticate.token": True,
+        "spark.connect.grpc.binding.port": None,
+        "spark.connect.grpc.arrow.maxBatchSize": None,
+        # Defaults from Spark https://github.com/apache/spark/blob/master/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveUtils.scala
+        "spark.sql.hive.version": "2.3.9",
+        "spark.sql.hive.metastore.version": "2.3.9",
+        "spark.sql.hive.metastore.jars": "builtin",
+        "spark.sql.hive.metastore.jars.path": None,
+        "spark.sql.hive.metastore.sharedPrefixes": [
+            "com.mysql.jdbc",
+            "com.mysql.cj",
+            "org.postgresql",
+            "com.microsoft.sqlserver",
+            "oracle.jdbc",
+        ],
+        "spark.sql.hive.metastore.barrierPrefixes": None,
+        "spark.sql.globalTempDatabase": "global_temp",
+    }
+
+    readonly_config_list = [
+        "snowpark.connect.version",
+    ]
+
+    default_global_config = {
+        # TODO: This will need to be changed to the actual host of the driver.
+        "spark.driver.host": "sc://localhost:15002",
+        "spark.sql.pyspark.inferNestedDictAsStruct.enabled": "false",
+        "spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled": "false",
+        "spark.sql.repl.eagerEval.enabled": "false",
+        "spark.sql.repl.eagerEval.maxNumRows": "20",
+        "spark.sql.repl.eagerEval.truncate": "20",
+        "spark.sql.session.localRelationCacheThreshold": "2147483647",
+        "spark.sql.session.timeZone": get_localzone_name(),  # spark by default uses jvm local timezone
+        "spark.sql.timestampType": "TIMESTAMP_LTZ",
+        "spark.sql.crossJoin.enabled": "true",
+        "spark.sql.caseSensitive": "false",
+        "spark.sql.mapKeyDedupPolicy": "EXCEPTION",
+        "spark.sql.ansi.enabled": "false",
+        "spark.sql.sources.default": "parquet",
+        "spark.Catalog.databaseFilterInformationSchema": "false",
+        "spark.sql.parser.quotedRegexColumnNames": "false",
+        # custom configs
+        "snowpark.connect.version": ".".join(map(str, sas_version)),
+    }
+
+    boolean_config_list = [
+        "spark.sql.pyspark.inferNestedDictAsStruct.enabled",
+        "spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled",
+        "spark.sql.repl.eagerEval.enabled",
+        "spark.sql.crossJoin.enabled",
+        "spark.sql.caseSensitive",
+        "spark.sql.ansi.enabled",
+        "spark.Catalog.databaseFilterInformationSchema",
+        "spark.sql.parser.quotedRegexColumnNames",
+    ]
+
+    int_config_list = [
+        "spark.sql.repl.eagerEval.maxNumRows",
+        "spark.sql.repl.eagerEval.truncate",
+        "spark.sql.session.localRelationCacheThreshold",
+    ]
+
+    # This is a mapping of the configuration key to a function that assigns the
+    # configuration value to a Snowpark session. This is only needed for configurations
+    # that we need to make available to the Snowpark session.
+    snowpark_config_mapping = {
+        # Snowpark uses query_tag for the application name.
+        "spark.app.name": lambda session, name: setattr(
+            session, "query_tag", f"Spark-Connect-App-Name={name}"
+        ),
+        "snowpark.connect.udf.packages": lambda session, packages: session.add_packages(
+            *packages.strip("[] ").split(",")
+        ),
+    }
+
+    float_config_list = []
+
+    def __init__(self) -> None:
+        self.global_config = SynchronizedDict(copy(self.default_global_config))
+        for key in self.global_config.keys():
+            setattr(self, key.replace(".", "_"), self._get_config_setting(key))
+
+    def _get_config_setting(self, key: str) -> bool | int | float | str | None:
+        """Get the configuration setting for the key based on the setting type."""
+        if key in self.boolean_config_list:
+            return str_to_bool(self.global_config[key])
+        elif key in self.int_config_list:
+            return int(self.global_config[key])
+        elif key in self.float_config_list:
+            return float(self.global_config[key])
+        else:
+            return self.global_config[key]
+
+    def __getattr__(self, item):
+        return self.get(item.replace("_", "."))
+
+    def __getitem__(self, item: str) -> str:
+        return self.get(item)
+
+    def __setitem__(self, key: str, value: str) -> None:
+        return self.set(key, value)
+
+    def get(self, key, default=None) -> str:
+        self._initialize_if_static_config_not_set(key)
+        return self.global_config.get(key, default)
+
+    def get_all(self) -> dict[str, str]:
+        return self.global_config.copy()
+
+    def set(self, key: str, value: str) -> None:
+        # Spark Connect sends us only string values, with empty string being
+        # equivalent to None.
+        if value == "":
+            value = None
+        self.global_config[key] = value
+        if key in self.snowpark_config_mapping.keys():
+            snowpark_session = get_or_create_snowpark_session()
+            self.snowpark_config_mapping[key](snowpark_session, value)
+        # This is necessary to make the configuration available as an attribute.
+        setattr(self, key.replace(".", "_"), self._get_config_setting(key))
+
+    def unset(self, key: str) -> None:
+        self.global_config.remove(key)
+
+    def is_modifiable(self, key) -> bool:
+        self._initialize_if_static_config_not_set(key)
+        is_in_config = self.is_set(key)
+        is_in_static_config = self.is_static_config(key)
+        return is_in_config and not is_in_static_config
+
+    def _initialize_if_static_config_not_set(self, key):
+        """
+        Spark maintains set of 'static' config values that are immutable.
+        SAS allows static configs only to be set once and only if they were never read before.
+        This function initializes static configs with default values if they haven't been set.
+        """
+        if self.is_static_config(key) and not self.is_set(key):
+            default_value = self.default_static_global_config[key]
+            self.set(key, default_value)
+            snowflake_session = get_or_create_snowpark_session()
+            set_snowflake_parameters(key, default_value, snowflake_session)
+
+    def is_static_config(self, key):
+        return key in self.default_static_global_config.keys()
+
+    def is_set(self, key):
+        return key in self.global_config.keys()
+
+
+SESSION_CONFIG_KEY_WHITELIST = {
+    "spark.hadoop.fs.s3a.access.key",
+    "spark.hadoop.fs.s3a.secret.key",
+    "spark.hadoop.fs.s3a.session.token",
+    "spark.sql.execution.pythonUDTF.arrow.enabled",
+    "spark.sql.tvf.allowMultipleTableArguments.enabled",
+    "snowpark.connect.sql.passthrough",
+    "snowpark.connect.iceberg.external_volume",
+    "snowpark.connect.auto-uppercase.ddl",
+    "snowpark.connect.auto-uppercase.dml",
+    "snowpark.connect.udtf.compatibility_mode",
+    "snowpark.connect.views.duplicate_column_names_handling_mode",
+}
+AZURE_SAS_KEY = re.compile(
+    r"^fs\.azure\.sas\.[^\.]+\.[^\.]+\.blob\.core\.windows\.net$"
+)
+
+
+def valid_session_config_key(key: str):
+    return (
+        key in SESSION_CONFIG_KEY_WHITELIST  # AWS session keys
+        or AZURE_SAS_KEY.match(key)  # Azure session keys
+    )
+
+
+class SessionConfig:
+    """This class contains the session configuration for the Spark Server."""
+
+    default_session_config = {
+        "snowpark.connect.auto-uppercase.ddl": "true",
+        "snowpark.connect.auto-uppercase.dml": "true",
+        "snowpark.connect.sql.passthrough": "false",
+        "snowpark.connect.udtf.compatibility_mode": "false",
+        "snowpark.connect.views.duplicate_column_names_handling_mode": "rename",
+        "spark.sql.execution.pythonUDTF.arrow.enabled": "false",
+        "spark.sql.tvf.allowMultipleTableArguments.enabled": "true",
+    }
+
+    def __init__(self) -> None:
+        self.config = copy(self.default_session_config)
+
+    def __getitem__(self, item: str) -> str:
+        return self.get(item)
+
+    def __setitem__(self, key: str, value: str) -> None:
+        return self.set(key, value)
+
+    def get(self, key, default="") -> str:
+        return self.config.get(key, default)
+
+    def set(self, key: str, value: str) -> None:
+        if not valid_session_config_key(key):
+            return
+
+        self.config[key] = value
+
+
+CONFIG_ALLOWED_VALUES: dict[str, tuple] = {
+    "snowpark.connect.views.duplicate_column_names_handling_mode": (
+        "rename",
+        "fail",
+        "drop",
+    )
+}
+
+# Set some default configuration that are necessary for the driver.
+global_config = GlobalConfig()
+sessions_config: defaultdict[str, SessionConfig] = defaultdict(SessionConfig)
+
+
+def route_config_proto(
+    config: proto_base.ConfigRequest,
+    snowpark_session: snowpark.Session,
+) -> proto_base.ConfigResponse:
+    global global_config
+    global sessions_config
+
+    op_type = config.operation.WhichOneof("op_type")
+    telemetry.report_config_op_type(op_type)
+    match op_type:
+        case "set":
+            logger.info("SET")
+
+            for pair in config.operation.set.pairs:
+                # Check if the value field is present, not present when invalid fields are set in conf.
+                if not pair.HasField("value"):
+                    from pyspark.errors import IllegalArgumentException
+
+                    raise IllegalArgumentException(
+                        f"Cannot set config '{pair.key}' to None"
+                    )
+
+                telemetry.report_config_set(pair.key, pair.value)
+                set_config_param(
+                    config.session_id, pair.key, pair.value, snowpark_session
+                )
+
+            return proto_base.ConfigResponse(session_id=config.session_id)
+        case "unset":
+            logger.info("UNSET")
+            for key in config.operation.unset.keys:
+                telemetry.report_config_unset(key)
+                unset_config_param(config.session_id, key, snowpark_session)
+
+            return proto_base.ConfigResponse(session_id=config.session_id)
+        case "get":
+            logger.info("GET")
+            res = proto_base.ConfigResponse(session_id=config.session_id)
+            for key in config.operation.get.keys:
+                pair = res.pairs.add()
+                pair.key = key
+                val = global_config.get(key)
+                if val:
+                    pair.value = str(val)
+            return res
+        case "get_with_default":
+            logger.info("GET_WITH_DEFAULT")
+            result_pairs = [
+                proto_base.KeyValue(
+                    key=pair.key,
+                    value=global_config.get(
+                        pair.key, pair.value if pair.HasField("value") else None
+                    ),
+                )
+                for pair in config.operation.get_with_default.pairs
+            ]
+            return proto_base.ConfigResponse(
+                session_id=config.session_id,
+                pairs=result_pairs,
+            )
+        case "get_option":
+            logger.info("GET_OPTION")
+            res = proto_base.ConfigResponse(session_id=config.session_id)
+            for key in config.operation.get_option.keys:
+                pair = res.pairs.add()
+                pair.key = key
+                val = global_config.get(key, None)
+                if val:
+                    pair.value = str(val)
+            return res
+        case "get_all":
+            logger.info("GET_ALL")
+            res = proto_base.ConfigResponse(session_id=config.session_id)
+            prefix = (
+                config.operation.get_all.prefix
+                if config.operation.get_all.HasField("prefix")
+                else None
+            )
+            config_items = global_config.get_all()
+            for item in config_items.items():
+                if prefix is None or item[0].startswith(prefix):
+                    pair = res.pairs.add()
+                    pair.key = (
+                        item[0] if prefix is None else item[0].removeprefix(prefix)
+                    )
+                    if item[1]:
+                        pair.value = str(item[1])
+            return res
+        case "is_modifiable":
+            logger.info("IS_MODIFIABLE")
+            res = proto_base.ConfigResponse(session_id=config.session_id)
+            for key in config.operation.is_modifiable.keys:
+                pair = res.pairs.add()
+                pair.key = key
+                # Change value to lowercase in order to mimic Java boolean type.
+                pair.value = str(global_config.is_modifiable(key)).lower()
+            return res
+        case _:
+            raise SnowparkConnectNotImplementedError(f"Unexpected request {config}")
+
+
+def set_config_param(
+    session_id: str, key, val, snowpark_session: snowpark.Session
+) -> None:
+    _verify_static_config_not_modified(key)
+    _verify_is_not_readonly_config(key)
+    _verify_is_valid_config_value(key, val)
+
+    global_config[key] = val
+    if valid_session_config_key(key):
+        sessions_config[session_id][key] = val
+    set_snowflake_parameters(key, val, snowpark_session)
+
+
+def unset_config_param(
+    session_id: str, key, snowpark_session: snowpark.Session
+) -> None:
+    _verify_static_config_not_modified(key)
+
+    default_value = global_config.default_global_config.get(key, None)
+
+    # Remove the key from global config
+    if default_value is None:
+        global_config.unset(key)
+    else:
+        global_config[key] = default_value
+
+    if valid_session_config_key(key):
+        default_session_value = SessionConfig.default_session_config.get(key, "")
+        sessions_config[session_id].set(key, default_session_value)
+
+    set_snowflake_parameters(key, default_value, snowpark_session)
+
+
+def _verify_static_config_not_modified(key: str) -> None:
+    # https://github.com/apache/spark/blob/v3.5.3/sql/core/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala#L161
+    # Spark does not allow to modify static configurations at runtime.
+    if global_config.is_static_config(key) and global_config.is_set(key):
+        raise ValueError(f"Cannot modify the value of a static config: {key}")
+
+
+def _verify_is_valid_config_value(key: str, value: Any) -> None:
+    if key in CONFIG_ALLOWED_VALUES and value not in CONFIG_ALLOWED_VALUES[key]:
+        raise ValueError(
+            f"Invalid value '{value}' for key '{key}'. Allowed values: {', '.join(CONFIG_ALLOWED_VALUES[key])}."
+        )
+
+
+def _verify_is_not_readonly_config(key):
+    if key in global_config.readonly_config_list:
+        raise ValueError(f"Config with key {key} is read-only and cannot be modified.")
+
+
+def set_jvm_timezone(timezone_id: str):
+    """
+    Set JVM default timezone at runtime (after JVM startup).
+
+    Args:
+        timezone_id: Timezone ID like 'America/New_York', 'UTC', 'Europe/London'
+
+    Returns:
+        bool: True if successfully set
+
+    Raises:
+        RuntimeError: If JVM is not started
+    """
+    if not jpype.isJVMStarted():
+        raise RuntimeError("JVM must be started before setting timezone")
+
+    try:
+        TimeZone = jpype.JClass("java.util.TimeZone")
+        new_timezone = TimeZone.getTimeZone(timezone_id)
+        TimeZone.setDefault(new_timezone)
+
+        logger.info(f"JVM timezone changed to: {new_timezone.getID()}")
+    except Exception as e:
+        logger.error(f"Failed to set JVM timezone: {e}")
+
+
+def reset_jvm_timezone_to_system_default():
+    """Reset JVM timezone to the system's default timezone"""
+    if not jpype.isJVMStarted():
+        raise RuntimeError("JVM must be started first")
+
+    try:
+        TimeZone = jpype.JClass("java.util.TimeZone")
+        TimeZone.setDefault(None)
+        logger.info(
+            f"Reset JVM timezone to system default: {TimeZone.getDefault().getID()}"
+        )
+    except jpype.JException as e:
+        raise RuntimeError(f"Java exception while resetting timezone: {e}")
+    except Exception as e:
+        raise RuntimeError(f"Unexpected error resetting JVM timezone: {e}")
+
+
+def set_snowflake_parameters(
+    key: str, value: Any, snowpark_session: snowpark.Session
+) -> None:
+    match key:
+        case "spark.sql.session.timeZone":
+            if value:
+                snowpark_session.sql(
+                    f"ALTER SESSION SET TIMEZONE = '{value}'"
+                ).collect()
+                set_jvm_timezone(value)
+                if hasattr(time, "tzset"):
+                    os.environ["TZ"] = value
+                    time.tzset()
+            else:
+                snowpark_session.sql("ALTER SESSION UNSET TIMEZONE").collect()
+                reset_jvm_timezone_to_system_default()
+                if hasattr(time, "tzset") and False:
+                    if "TZ" in os.environ:
+                        del os.environ["TZ"]
+                    time.tzset()
+        case "spark.sql.globalTempDatabase":
+            if not value:
+                value = global_config.default_static_global_config.get(key)
+
+            snowpark_name = quote_name_without_upper_casing(value)
+            if auto_uppercase_ddl():
+                snowpark_name = snowpark_name.upper()
+
+            # Create the schema on demand. Before creating it, however,
+            # check if it already exists, to handle the case where the schema exists
+            # but the user does not have permissions to execute `CREATE SCHEMA`.
+            try:
+                snowpark_session.sql(
+                    "DESC SCHEMA identifier(?)", [snowpark_name]
+                ).collect()
+            except SnowparkSQLException:
+                db = snowpark_session.get_current_database()
+                prev_schema = snowpark_session.get_current_schema()
+                # Even though we checked that the schema doesn't exist,
+                # use "IF NOT EXISTS" to avoid race conditions.
+                snowpark_session.sql(
+                    "CREATE SCHEMA IF NOT EXISTS identifier(?)", [snowpark_name]
+                ).collect()
+                # When schema is created, it is also changed in the context of the session
+                match (prev_schema, snowpark_session.get_current_schema()):
+                    case (None, curr) if curr is not None:
+                        # session.use_schema(None) is not allowed.
+                        # Instead, we call `use_database` which will set schema in context to the default one.
+                        snowpark_session.use_database(db)
+                    case (prev, curr) if prev != curr:
+                        snowpark_session.use_schema(prev)
+        case _:
+            pass
+
+
+def get_boolean_session_config_param(name: str) -> bool:
+    session_config = sessions_config[get_session_id()]
+    return str_to_bool(session_config[name])
+
+
+def auto_uppercase_dml() -> bool:
+    return get_boolean_session_config_param("snowpark.connect.auto-uppercase.dml")
+
+
+def auto_uppercase_ddl() -> bool:
+    return get_boolean_session_config_param("snowpark.connect.auto-uppercase.ddl")

snowflake/snowpark_connect/constants.py

@@ -0,0 +1,47 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+DEFAULT_CONNECTION_NAME = "spark-connect"
+DEFAULT_CONNECTION_NAME_IN_SPCS = "default"
+DEFAULT_SNOWPARK_SUBMIT_CONNECTION_NAME = "snowpark-submit"
+
+SERVER_SIDE_SESSION_ID = "321"
+
+STRUCTURED_TYPES_ENABLED = True
+
+# UDF evaluation types
+MAP_IN_ARROW_EVAL_TYPE = 207  # eval_type for mapInArrow operations
+
+SPARK_TZ_ABBREVIATIONS_OVERRIDES = {
+    "ACT": "Australia/Darwin",
+    "AET": "Australia/Sydney",
+    "AGT": "America/Argentina/Buenos_Aires",
+    "ART": "Africa/Cairo",
+    "AST": "America/Anchorage",
+    "BET": "America/Sao_Paulo",
+    "BST": "Asia/Dhaka",
+    "CAT": "Africa/Harare",
+    "CNT": "America/St_Johns",
+    "CST": "America/Chicago",
+    "CTT": "Asia/Shanghai",
+    "EAT": "Africa/Addis_Ababa",
+    "ECT": "Europe/Paris",
+    "IET": "America/Indiana/Indianapolis",
+    "IST": "Asia/Kolkata",
+    "JST": "Asia/Tokyo",
+    "MIT": "Pacific/Apia",
+    "NET": "Asia/Yerevan",
+    "NST": "Pacific/Auckland",
+    "PLT": "Asia/Karachi",
+    "PNT": "America/Phoenix",
+    "PRT": "America/Puerto_Rico",
+    "PST": "America/Los_Angeles",
+    "SST": "Pacific/Guadalcanal",
+    "VST": "Asia/Ho_Chi_Minh",
+}
+
+
+COLUMN_METADATA_COLLISION_KEY = "{expr_id}_{key}"
+
+DUPLICATE_KEY_FOUND_ERROR_TEMPLATE = "Duplicate key found: {key}. You can set spark.sql.mapKeyDedupPolicy to LAST_WIN to deduplicate map keys with last wins policy."

snowflake/snowpark_connect/control_server.py

@@ -0,0 +1,52 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import snowflake.snowpark_connect.proto.control_pb2 as control_proto
+import snowflake.snowpark_connect.proto.control_pb2_grpc as control_grpc
+from snowflake import snowpark
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+
+
+class ControlServicer(control_grpc.ControlServiceServicer):
+    def __init__(self, session: snowpark.session.Session) -> None:
+        self.session = session
+        self.log_ast = True  # TODO: False after Configure is enabled.
+        self.ast_listener = session.ast_listener()
+        self.spark_connect_batches = []
+
+    def log_spark_connect_batch(self, batch):
+        self.spark_connect_batches.append(batch)
+
+    # TODO: Enable configure method when it is ready.
+    # def Configure(self, request: control_proto.Config, context):
+    #     if request.HasField("log_ast"):
+    #         self.log_ast = request.log_ast
+    #         # TODO: attach/detach listener.
+    #
+    #     return Config(log_ast=self.log_ast)
+
+    def Ping(
+        self, request: control_proto.PingRequest, context
+    ) -> control_proto.PingResponse:
+        return control_proto.PingResponse(payload=request.payload + " **")
+
+    def GetRequestAst(
+        self, request: control_proto.GetRequestAstRequest, context
+    ) -> control_proto.GetRequestAstResponse:
+        # TODO: assert self.session.ast_enabled
+        logger.info(
+            f"Returning AST for {len(self.ast_listener._ast_batches)} batches and "
+            + f"{len(self.spark_connect_batches)} Spark requests"
+        )
+        res = control_proto.GetRequestAstResponse()
+
+        # Copy the batches.
+        res.spark_requests.extend(self.spark_connect_batches)
+        res.snowpark_ast_batches.extend(self.ast_listener._ast_batches)
+
+        # Reset the local state.
+        self.spark_connect_batches = []
+        self.ast_listener._ast_batches = []
+
+        return res