Jun 16, 2025 · Apr 16, 2025 · May 12, 2025 · May 17, 2025 · May 17, 2025 · May 19, 2025
diff --git a/src/databricks/sql/parameters/__init__.py b/src/databricks/sql/parameters/__init__.py
    TimestampNTZParameter,
    TinyIntParameter,
    DecimalParameter,
    MapParameter,
    ArrayParameter,
 )
diff --git a/src/databricks/sql/parameters/native.py b/src/databricks/sql/parameters/native.py
 import datetime
 import decimal
 from enum import Enum, auto
 from typing import Optional, Sequence
 from typing import Optional, Sequence, Any

 from databricks.sql.exc import NotSupportedError
 from databricks.sql.thrift_api.TCLIService.ttypes import (
    TSparkParameter,
    TSparkParameterValue,
    TSparkParameterValueArg,
 )

 import datetime


 TAllowedParameterValue = Union[
    str, int, float, datetime.datetime, datetime.date, bool, decimal.Decimal, None
    str,
    int,
    float,
    datetime.datetime,
    datetime.date,
    bool,
    decimal.Decimal,
    None,
    list,
    dict,
    tuple,
 ]



    CAST_EXPR: str
    name: Optional[str]
    value: Any

    def as_tspark_param(self, named: bool) -> TSparkParameter:
        """Returns a TSparkParameter object that can be passed to the DBR thrift server."""
    def _tspark_param_value(self):
        return TSparkParameterValue(stringValue=str(self.value))

    def _tspark_value_arg(self):
        """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
        return TSparkParameterValueArg(value=str(self.value), type=self._cast_expr())

    def _cast_expr(self):
        return self.CAST_EXPR

    CAST_EXPR = DatabricksSupportedType.TINYINT.name


 class ArrayParameter(DbsqlParameterBase):
    """Wrap a Python `Sequence` that will be bound to a Databricks SQL ARRAY type."""

    def __init__(self, value: Sequence[Any], name: Optional[str] = None):
        """
        :value:
            The value to bind for this parameter. This will be casted to a ARRAY.
        :name:
            If None, your query must contain a `?` marker. Like:

            ```sql
               SELECT * FROM table WHERE field = ?
            ```
            If not None, your query should contain a named parameter marker. Like:
            ```sql
                SELECT * FROM table WHERE field = :my_param
            ```

            The `name` argument to this function would be `my_param`.
        """
        self.name = name
        self.value = [dbsql_parameter_from_primitive(val) for val in value]

    def as_tspark_param(self, named: bool = False) -> TSparkParameter:
        """Returns a TSparkParameter object that can be passed to the DBR thrift server."""

        tsp = TSparkParameter(type=self._cast_expr())
        tsp.arguments = [val._tspark_value_arg() for val in self.value]

        if named:
            tsp.name = self.name
            tsp.ordinal = False
        elif not named:
            tsp.ordinal = True
        return tsp

    def _tspark_value_arg(self):
        """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
        tva = TSparkParameterValueArg(type=self._cast_expr())
        tva.arguments = [val._tspark_value_arg() for val in self.value]
        return tva

    CAST_EXPR = DatabricksSupportedType.ARRAY.name


 class MapParameter(DbsqlParameterBase):
    """Wrap a Python `dict` that will be bound to a Databricks SQL MAP type."""

    def __init__(self, value: dict, name: Optional[str] = None):
        """
        :value:
            The value to bind for this parameter. This will be casted to a MAP.
        :name:
            If None, your query must contain a `?` marker. Like:

            ```sql
               SELECT * FROM table WHERE field = ?
            ```
            If not None, your query should contain a named parameter marker. Like:
            ```sql
                SELECT * FROM table WHERE field = :my_param
            ```

            The `name` argument to this function would be `my_param`.
        """
        self.name = name
        self.value = [
            dbsql_parameter_from_primitive(item)
            for key, val in value.items()
            for item in (key, val)
        ]

    def as_tspark_param(self, named: bool = False) -> TSparkParameter:
        """Returns a TSparkParameter object that can be passed to the DBR thrift server."""

        tsp = TSparkParameter(type=self._cast_expr())
        tsp.arguments = [val._tspark_value_arg() for val in self.value]
        if named:
            tsp.name = self.name
            tsp.ordinal = False
        elif not named:
            tsp.ordinal = True
        return tsp

    def _tspark_value_arg(self):
        """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
        tva = TSparkParameterValueArg(type=self._cast_expr())
        tva.arguments = [val._tspark_value_arg() for val in self.value]
        return tva

    CAST_EXPR = DatabricksSupportedType.MAP.name


 class DecimalParameter(DbsqlParameterBase):
    """Wrap a Python `Decimal` that will be bound to a Databricks SQL DECIMAL type."""

    # havoc. We can't use TYPE_INFERRENCE_MAP because mypy doesn't trust
    # its logic

    if type(value) is int:
    if isinstance(value, bool):
        return BooleanParameter(value=value, name=name)
    elif isinstance(value, int):
        return dbsql_parameter_from_int(value, name=name)
    eliftype(value) isstr:
    elifisinstance(value,str):
        return StringParameter(value=value, name=name)
    eliftype(value) isfloat:
    elifisinstance(value,float):
        return FloatParameter(value=value, name=name)
    eliftype(value) isdatetime.datetime:
    elifisinstance(value,datetime.datetime):
        return TimestampParameter(value=value, name=name)
    eliftype(value) isdatetime.date:
    elifisinstance(value,datetime.date):
        return DateParameter(value=value, name=name)
    elif type(value) is bool:
        return BooleanParameter(value=value, name=name)
    elif type(value) is decimal.Decimal:
    elif isinstance(value, decimal.Decimal):
        return DecimalParameter(value=value, name=name)
    elif isinstance(value, dict):
        return MapParameter(value=value, name=name)
    elif isinstance(value, Sequence) and not isinstance(value, str):
        return ArrayParameter(value=value, name=name)
    elif value is None:
        return VoidParameter(value=value, name=name)

    else:
        raise NotSupportedError(
            f"Could not infer parameter type from value: {value} - {type(value)} \n"
    TimestampNTZParameter,
    TinyIntParameter,
    DecimalParameter,
    ArrayParameter,
    MapParameter,
 ]


diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
 import decimal
 from abc import ABC, abstractmethod
 from collections import OrderedDict, namedtuple
 from collections.abc importIterable
 from collections.abc importMapping
 from decimal import Decimal
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 from typing import Any, Dict, List, Optional, Union, Sequence
 import re

 import lz4.frame
 # Taken from PyHive
 class ParamEscaper:
    _DATE_FORMAT = "%Y-%m-%d"
    _TIME_FORMAT = "%H:%M:%S.%f"
    _TIME_FORMAT = "%H:%M:%S.%f %z"
    _DATETIME_FORMAT = "{} {}".format(_DATE_FORMAT, _TIME_FORMAT)

    def escape_args(self, parameters):
        return "'{}'".format(item.replace("\\", "\\\\").replace("'", "\\'"))

    def escape_sequence(self, item):
        l = map(str, map(self.escape_item, item))
        return "(" + ",".join(l) + ")"
        l = map(self.escape_item, item)
        l = list(map(str, l))
        return "ARRAY(" + ",".join(l) + ")"

    def escape_mapping(self, item):
        l = map(
            self.escape_item,
            (element for key, value in item.items() for element in (key, value)),
        )
        l = list(map(str, l))
        return "MAP(" + ",".join(l) + ")"

    def escape_datetime(self, item, format, cutoff=0):
        dt_str = item.strftime(format)
        formatted = dt_str[:-cutoff] if cutoff and format.endswith(".%f") else dt_str
        return "'{}'".format(formatted)
        return "'{}'".format(formatted.strip())

    def escape_decimal(self, item):
        return str(item)
            return self.escape_number(item)
        elif isinstance(item, str):
            return self.escape_string(item)
        elif isinstance(item, Iterable):
            return self.escape_sequence(item)
        elif isinstance(item, datetime.datetime):
            return self.escape_datetime(item, self._DATETIME_FORMAT)
        elif isinstance(item, datetime.date):
            return self.escape_datetime(item, self._DATE_FORMAT)
        elif isinstance(item, decimal.Decimal):
            return self.escape_decimal(item)
        elif isinstance(item, Sequence):
            return self.escape_sequence(item)
        elif isinstance(item, Mapping):
            return self.escape_mapping(item)
        else:
            raise exc.ProgrammingError("Unsupported object {}".format(item))

diff --git a/tests/e2e/test_complex_types.py b/tests/e2e/test_complex_types.py
 import pytest
 from numpy import ndarray
 from typing import Sequence

 from tests.e2e.test_driver import PySQLPytestTestCase

            # Create the table
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTSpysql_e2e_test_complex_types_table (
                CREATE TABLE IF NOT EXISTSpysql_test_complex_types_table (
                    array_col ARRAY<STRING>,
                    map_col MAP<STRING, INTEGER>,
                    struct_col STRUCT<field1: STRING, field2: INTEGER>
                )
                    struct_col STRUCT<field1: STRING, field2: INTEGER>,
                    array_array_col ARRAY<ARRAY<STRING>>,
                    array_map_col ARRAY<MAP<STRING, INTEGER>>,
                    map_array_col MAP<STRING, ARRAY<STRING>>
                ) USING DELTA
                """
            )
            # Insert a record
            cursor.execute(
                """
                INSERT INTOpysql_e2e_test_complex_types_table
                INSERT INTOpysql_test_complex_types_table
                VALUES (
                    ARRAY('a', 'b', 'c'),
                    MAP('a', 1, 'b', 2, 'c', 3),
                    NAMED_STRUCT('field1', 'a', 'field2', 1)
                    NAMED_STRUCT('field1', 'a', 'field2', 1),
                    ARRAY(ARRAY('a','b','c')),
                    ARRAY(MAP('a', 1, 'b', 2, 'c', 3)),
                    MAP('a', ARRAY('a', 'b', 'c'), 'b', ARRAY('d', 'e'))
                )
                """
            )
            yield
            # Clean up the table after the test
            cursor.execute("DROP TABLE IF EXISTS pysql_e2e_test_complex_types_table")
            cursor.execute("DELETE FROM pysql_test_complex_types_table")

    @pytest.mark.parametrize(
        "field,expected_type",
        [("array_col", ndarray), ("map_col", list), ("struct_col", dict)],
        [
            ("array_col", ndarray),
            ("map_col", list),
            ("struct_col", dict),
            ("array_array_col", ndarray),
            ("array_map_col", ndarray),
            ("map_array_col", list),
        ],
    )
    def test_read_complex_types_as_arrow(self, field, expected_type, table_fixture):
        """Confirms the return types of a complex type field when reading as arrow"""

        with self.cursor() as cursor:
            result = cursor.execute(
                "SELECT * FROMpysql_e2e_test_complex_types_table LIMIT 1"
                "SELECT * FROMpysql_test_complex_types_table LIMIT 1"
            ).fetchone()

        assert isinstance(result[field], expected_type)

    @pytest.mark.parametrize("field", [("array_col"), ("map_col"), ("struct_col")])
    @pytest.mark.parametrize(
        "field",
        [
            ("array_col"),
            ("map_col"),
            ("struct_col"),
            ("array_array_col"),
            ("array_map_col"),
            ("map_array_col"),
        ],
    )
    def test_read_complex_types_as_string(self, field, table_fixture):
        """Confirms the return type of a complex type that is returned as a string"""
        with self.cursor(
            extra_params={"_use_arrow_native_complex_types": False}
        ) as cursor:
            result = cursor.execute(
                "SELECT * FROMpysql_e2e_test_complex_types_table LIMIT 1"
                "SELECT * FROMpysql_test_complex_types_table LIMIT 1"
            ).fetchone()

        assert isinstance(result[field], str)
diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py
                    raise KeyboardInterrupt("Simulated interrupt")
        finally:
            if conn is not None:
                assert not conn.open, "Connection should be closed after KeyboardInterrupt"
                assert (
                    not conn.open
                ), "Connection should be closed after KeyboardInterrupt"

    def test_cursor_close_properly_closes_operation(self):
        """Test that Cursor.close() properly closes the active operation handle on the server."""
                        raise KeyboardInterrupt("Simulated interrupt")
        finally:
            if cursor is not None:
                assert not cursor.open, "Cursor should be closed after KeyboardInterrupt"
                assert (
                    not cursor.open
                ), "Cursor should be closed after KeyboardInterrupt"

    def test_nested_cursor_context_managers(self):
        """Test that nested cursor context managers properly close operations on the server."""
Original file line number	Diff line number	Diff line change
Expand Up		@@ -12,4 +12,6 @@
		TimestampNTZParameter,
		TinyIntParameter,
		DecimalParameter,
		MapParameter,
		ArrayParameter,
		)
Original file line number	Diff line number	Diff line change
		@@ -1,12 +1,13 @@
		import datetime
		import decimal
		from enum import Enum, auto
		from typing import Optional, Sequence
		from typing import Optional, Sequence, Any

		from databricks.sql.exc import NotSupportedError
		from databricks.sql.thrift_api.TCLIService.ttypes import (
		TSparkParameter,
		TSparkParameterValue,
		TSparkParameterValueArg,
		)

		import datetime
Expand DownExpand Up		@@ -54,7 +55,17 @@ class DatabricksSupportedType(Enum):


		TAllowedParameterValue = Union[
		str, int, float, datetime.datetime, datetime.date, bool, decimal.Decimal, None
		str,
		int,
		float,
		datetime.datetime,
		datetime.date,
		bool,
		decimal.Decimal,
		None,
		list,
		dict,
		tuple,
		]


Expand DownExpand Up		@@ -82,6 +93,7 @@ class DbsqlParameterBase:

		CAST_EXPR: str
		name: Optional[str]
		value: Any

		def as_tspark_param(self, named: bool) -> TSparkParameter:
		"""Returns a TSparkParameter object that can be passed to the DBR thrift server."""
Expand All		@@ -98,6 +110,10 @@ def as_tspark_param(self, named: bool) -> TSparkParameter:
		def _tspark_param_value(self):
		return TSparkParameterValue(stringValue=str(self.value))

		def _tspark_value_arg(self):
		"""Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
		return TSparkParameterValueArg(value=str(self.value), type=self._cast_expr())

		def _cast_expr(self):
		return self.CAST_EXPR

Expand DownExpand Up		@@ -428,6 +444,99 @@ def __init__(self, value: int, name: Optional[str] = None):
		CAST_EXPR = DatabricksSupportedType.TINYINT.name


		class ArrayParameter(DbsqlParameterBase):
		"""Wrap a Python `Sequence` that will be bound to a Databricks SQL ARRAY type."""

		def __init__(self, value: Sequence[Any], name: Optional[str] = None):
		"""
		:value:
		The value to bind for this parameter. This will be casted to a ARRAY.
		:name:
		If None, your query must contain a `?` marker. Like:

		```sql
		SELECT * FROM table WHERE field = ?
		```
		If not None, your query should contain a named parameter marker. Like:
		```sql
		SELECT * FROM table WHERE field = :my_param
		```

		The `name` argument to this function would be `my_param`.
		"""
		self.name = name
		self.value = [dbsql_parameter_from_primitive(val) for val in value]
jprakash-db marked this conversation as resolved. Show resolvedHide resolved

		def as_tspark_param(self, named: bool = False) -> TSparkParameter:
		"""Returns a TSparkParameter object that can be passed to the DBR thrift server."""

		tsp = TSparkParameter(type=self._cast_expr())
		tsp.arguments = [val._tspark_value_arg() for val in self.value]

		if named:
		tsp.name = self.name
		tsp.ordinal = False
		elif not named:
		tsp.ordinal = True
jprakash-db marked this conversation as resolved. Show resolvedHide resolved
		return tsp

		def _tspark_value_arg(self):
		"""Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
		tva = TSparkParameterValueArg(type=self._cast_expr())
		tva.arguments = [val._tspark_value_arg() for val in self.value]
		return tva

		CAST_EXPR = DatabricksSupportedType.ARRAY.name


		class MapParameter(DbsqlParameterBase):
		"""Wrap a Python `dict` that will be bound to a Databricks SQL MAP type."""

		def __init__(self, value: dict, name: Optional[str] = None):
		"""
		:value:
		The value to bind for this parameter. This will be casted to a MAP.
		:name:
		If None, your query must contain a `?` marker. Like:

		```sql
		SELECT * FROM table WHERE field = ?
		```
		If not None, your query should contain a named parameter marker. Like:
		```sql
		SELECT * FROM table WHERE field = :my_param
		```

		The `name` argument to this function would be `my_param`.
		"""
		self.name = name
		self.value = [
		dbsql_parameter_from_primitive(item)
		for key, val in value.items()
		for item in (key, val)
		]

		def as_tspark_param(self, named: bool = False) -> TSparkParameter:
		"""Returns a TSparkParameter object that can be passed to the DBR thrift server."""

		tsp = TSparkParameter(type=self._cast_expr())
		tsp.arguments = [val._tspark_value_arg() for val in self.value]
		if named:
		tsp.name = self.name
		tsp.ordinal = False
		elif not named:
		tsp.ordinal = True
		return tsp

		def _tspark_value_arg(self):
		"""Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server."""
		tva = TSparkParameterValueArg(type=self._cast_expr())
		tva.arguments = [val._tspark_value_arg() for val in self.value]
		return tva

		CAST_EXPR = DatabricksSupportedType.MAP.name


		class DecimalParameter(DbsqlParameterBase):
		"""Wrap a Python `Decimal` that will be bound to a Databricks SQL DECIMAL type."""

Expand DownExpand Up		@@ -543,23 +652,26 @@ def dbsql_parameter_from_primitive(
		# havoc. We can't use TYPE_INFERRENCE_MAP because mypy doesn't trust
		# its logic

		if type(value) is int:
		if isinstance(value, bool):
		return BooleanParameter(value=value, name=name)
		elif isinstance(value, int):
		return dbsql_parameter_from_int(value, name=name)
		eliftype(value) isstr:
		elifisinstance(value,str):
		return StringParameter(value=value, name=name)
		eliftype(value) isfloat:
		elifisinstance(value,float):
		return FloatParameter(value=value, name=name)
		eliftype(value) isdatetime.datetime:
		elifisinstance(value,datetime.datetime):
		return TimestampParameter(value=value, name=name)
		eliftype(value) isdatetime.date:
		elifisinstance(value,datetime.date):
		return DateParameter(value=value, name=name)
		elif type(value) is bool:
		return BooleanParameter(value=value, name=name)
		elif type(value) is decimal.Decimal:
		elif isinstance(value, decimal.Decimal):
		return DecimalParameter(value=value, name=name)
		elif isinstance(value, dict):
		return MapParameter(value=value, name=name)
		elif isinstance(value, Sequence) and not isinstance(value, str):
jprakash-db marked this conversation as resolved. Show resolvedHide resolved
		return ArrayParameter(value=value, name=name)
		elif value is None:
		return VoidParameter(value=value, name=name)

		else:
		raise NotSupportedError(
		f"Could not infer parameter type from value: {value} - {type(value)} \n"
Expand All		@@ -581,6 +693,8 @@ def dbsql_parameter_from_primitive(
		TimestampNTZParameter,
		TinyIntParameter,
		DecimalParameter,
		ArrayParameter,
		MapParameter,
		]


Expand Down
Original file line number	Diff line number	Diff line change
		@@ -1,5 +1,6 @@
		import pytest
		from numpy import ndarray
		from typing import Sequence

		from tests.e2e.test_driver import PySQLPytestTestCase

Expand All		@@ -14,50 +15,73 @@ def table_fixture(self, connection_details):
		# Create the table
		cursor.execute(
		"""
		CREATE TABLE IF NOT EXISTSpysql_e2e_test_complex_types_table (
		CREATE TABLE IF NOT EXISTSpysql_test_complex_types_table (
		array_col ARRAY<STRING>,
		map_col MAP<STRING, INTEGER>,
		struct_col STRUCT<field1: STRING, field2: INTEGER>
		)
		struct_col STRUCT<field1: STRING, field2: INTEGER>,
		array_array_col ARRAY<ARRAY<STRING>>,
		array_map_col ARRAY<MAP<STRING, INTEGER>>,
		map_array_col MAP<STRING, ARRAY<STRING>>
		) USING DELTA
		"""
		)
		# Insert a record
		cursor.execute(
		"""
		INSERT INTOpysql_e2e_test_complex_types_table
		INSERT INTOpysql_test_complex_types_table
		VALUES (
		ARRAY('a', 'b', 'c'),
		MAP('a', 1, 'b', 2, 'c', 3),
		NAMED_STRUCT('field1', 'a', 'field2', 1)
		NAMED_STRUCT('field1', 'a', 'field2', 1),
		ARRAY(ARRAY('a','b','c')),
		ARRAY(MAP('a', 1, 'b', 2, 'c', 3)),
		MAP('a', ARRAY('a', 'b', 'c'), 'b', ARRAY('d', 'e'))
		)
		"""
		)
		yield
		# Clean up the table after the test
		cursor.execute("DROP TABLE IF EXISTS pysql_e2e_test_complex_types_table")
		cursor.execute("DELETE FROM pysql_test_complex_types_table")

		@pytest.mark.parametrize(
		"field,expected_type",
		[("array_col", ndarray), ("map_col", list), ("struct_col", dict)],
		[
		("array_col", ndarray),
		("map_col", list),
		("struct_col", dict),
		("array_array_col", ndarray),
		("array_map_col", ndarray),
		("map_array_col", list),
		],
		)
		def test_read_complex_types_as_arrow(self, field, expected_type, table_fixture):
		"""Confirms the return types of a complex type field when reading as arrow"""

		with self.cursor() as cursor:
		result = cursor.execute(
		"SELECT * FROMpysql_e2e_test_complex_types_table LIMIT 1"
		"SELECT * FROMpysql_test_complex_types_table LIMIT 1"
		).fetchone()

		assert isinstance(result[field], expected_type)

		@pytest.mark.parametrize("field", [("array_col"), ("map_col"), ("struct_col")])
		@pytest.mark.parametrize(
		"field",
		[
		("array_col"),
		("map_col"),
		("struct_col"),
		("array_array_col"),
		("array_map_col"),
		("map_array_col"),
		],
		)
		def test_read_complex_types_as_string(self, field, table_fixture):
		"""Confirms the return type of a complex type that is returned as a string"""
		with self.cursor(
		extra_params={"_use_arrow_native_complex_types": False}
		) as cursor:
		result = cursor.execute(
		"SELECT * FROMpysql_e2e_test_complex_types_table LIMIT 1"
		"SELECT * FROMpysql_test_complex_types_table LIMIT 1"
		).fetchone()

		assert isinstance(result[field], str)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -856,7 +856,9 @@ def test_closing_a_closed_connection_doesnt_fail(self, caplog):
		raise KeyboardInterrupt("Simulated interrupt")
		finally:
		if conn is not None:
		assert not conn.open, "Connection should be closed after KeyboardInterrupt"
		assert (
		not conn.open
		), "Connection should be closed after KeyboardInterrupt"

		def test_cursor_close_properly_closes_operation(self):
		"""Test that Cursor.close() properly closes the active operation handle on the server."""
Expand All		@@ -883,7 +885,9 @@ def test_cursor_close_properly_closes_operation(self):
		raise KeyboardInterrupt("Simulated interrupt")
		finally:
		if cursor is not None:
		assert not cursor.open, "Cursor should be closed after KeyboardInterrupt"
		assert (
		not cursor.open
		), "Cursor should be closed after KeyboardInterrupt"

		def test_nested_cursor_context_managers(self):
		"""Test that nested cursor context managers properly close operations on the server."""
Expand Down