[PECO-1026] Add Parameterized Query support to Python (#217)

* Initial commit

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Added tsparkparam handling

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Added basic test

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Addressed comments

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Addressed missed comments

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Resolved comments

---------

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

Author: nithinkdb

src/databricks/sql/client.py

@@ -14,7 +14,11 @@
     CursorAlreadyClosedError,
 )
 from databricks.sql.thrift_backend import ThriftBackend
-from databricks.sql.utils import ExecuteResponse, ParamEscaper, inject_parameters
+from databricks.sql.utils import (
+    ExecuteResponse,
+    ParamEscaper,
+    named_parameters_to_tsparkparams,
+)
 from databricks.sql.types import Row
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
 from databricks.sql.experimental.oauth_persistence import OAuthPersistence
@@ -482,7 +486,9 @@ def _handle_staging_remove(self, presigned_url: str, headers: dict = None):
             )
 
     def execute(
-        self, operation: str, parameters: Optional[Dict[str, str]] = None
+        self,
+        operation: str,
+        parameters: Optional[Union[List[Any], Dict[str, str]]] = None,
     ) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
@@ -493,10 +499,10 @@ def execute(
             Will result in the query "SELECT * FROM table WHERE field = 'foo' being sent to the server
         :returns self
         """
-        if parameters is not None:
-            operation = inject_parameters(
-                operation, self.escaper.escape_args(parameters)
-            )
+        if parameters is None:
+            parameters = []
+        else:
+            parameters = named_parameters_to_tsparkparams(parameters)
 
         self._check_not_closed()
         self._close_and_clear_active_result_set()
@@ -508,6 +514,7 @@ def execute(
             lz4_compression=self.connection.lz4_compression,
             cursor=self,
             use_cloud_fetch=self.connection.use_cloud_fetch,
+            parameters=parameters,
         )
         self.active_result_set = ResultSet(
             self.connection,

src/databricks/sql/thrift_backend.py

@@ -224,7 +224,7 @@ def __init__(
     def _initialize_retry_args(self, kwargs):
         # Configure retries & timing: use user-settings or defaults, and bound
         # by policy. Log.warn when given param gets restricted.
-        for (key, (type_, default, min, max)) in _retry_policy.items():
+        for key, (type_, default, min, max) in _retry_policy.items():
             given_or_default = type_(kwargs.get(key, default))
             bound = _bound(min, max, given_or_default)
             setattr(self, key, bound)
@@ -368,7 +368,6 @@ def attempt_request(attempt):
 
             error, error_message, retry_delay = None, None, None
             try:
-
                 this_method_name = getattr(method, "__name__")
 
                 logger.debug("Sending request: {}(<REDACTED>)".format(this_method_name))
@@ -614,7 +613,10 @@ def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, descripti
                 num_rows,
             ) = convert_column_based_set_to_arrow_table(t_row_set.columns, description)
         elif t_row_set.arrowBatches is not None:
-            (arrow_table, num_rows,) = convert_arrow_based_set_to_arrow_table(
+            (
+                arrow_table,
+                num_rows,
+            ) = convert_arrow_based_set_to_arrow_table(
                 t_row_set.arrowBatches, lz4_compressed, schema_bytes
             )
         else:
@@ -813,6 +815,7 @@ def execute_command(
         lz4_compression,
         cursor,
         use_cloud_fetch=False,
+        parameters=[],
     ):
         assert session_handle is not None
 
@@ -839,6 +842,7 @@ def execute_command(
                 "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
             },
             useArrowNativeTypes=spark_arrow_types,
+            parameters=parameters,
         )
         resp = self.make_request(self._client.ExecuteStatement, req)
         return self._handle_execute_response(resp, cursor)

src/databricks/sql/utils.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from abc import ABC, abstractmethod
 from collections import namedtuple, OrderedDict
 from collections.abc import Iterable
@@ -8,13 +9,17 @@
 import lz4.frame
 from typing import Dict, List, Union, Any
 import pyarrow
+from enum import Enum
+import copy
 
 from databricks.sql import exc, OperationalError
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TSparkArrowResultLink,
     TSparkRowSetType,
     TRowSet,
+    TSparkParameter,
+    TSparkParameterValue,
 )
 
 BIT_MASKS = [1, 2, 4, 8, 16, 32, 64, 128]
@@ -404,7 +409,7 @@ def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema
 
 
 def convert_decimals_in_arrow_table(table, description) -> pyarrow.Table:
-    for (i, col) in enumerate(table.itercolumns()):
+    for i, col in enumerate(table.itercolumns()):
         if description[i][1] == "decimal":
             decimal_col = col.to_pandas().apply(
                 lambda v: v if v is None else Decimal(v)
@@ -470,3 +475,86 @@ def _create_arrow_array(t_col_value_wrapper, arrow_type):
             result[i] = None
 
     return pyarrow.array(result, type=arrow_type)
+
+
+class DbSqlType(Enum):
+    STRING = "STRING"
+    DATE = "DATE"
+    TIMESTAMP = "TIMESTAMP"
+    FLOAT = "FLOAT"
+    DECIMAL = "DECIMAL"
+    INTEGER = "INTEGER"
+    BIGINT = "BIGINT"
+    SMALLINT = "SMALLINT"
+    TINYINT = "TINYINT"
+    BOOLEAN = "BOOLEAN"
+    INTERVAL_MONTH = "INTERVAL MONTH"
+    INTERVAL_DAY = "INTERVAL DAY"
+
+
+class DbSqlParameter:
+    name: str
+    value: Any
+    type: DbSqlType
+
+    def __init__(self, name="", value=None, type=None):
+        self.name = name
+        self.value = value
+        self.type = type
+
+    def __eq__(self, other):
+        return isinstance(other, self.__class__) and self.__dict__ == other.__dict__
+
+
+def named_parameters_to_dbsqlparams_v1(parameters: Dict[str, str]):
+    dbsqlparams = []
+    for name, parameter in parameters.items():
+        dbsqlparams.append(DbSqlParameter(name=name, value=parameter))
+    return dbsqlparams
+
+
+def named_parameters_to_dbsqlparams_v2(parameters: List[Any]):
+    dbsqlparams = []
+    for parameter in parameters:
+        if isinstance(parameter, DbSqlParameter):
+            dbsqlparams.append(parameter)
+        else:
+            dbsqlparams.append(DbSqlParameter(value=parameter))
+    return dbsqlparams
+
+
+def infer_types(params: list[DbSqlParameter]):
+    type_lookup_table = {
+        str: DbSqlType.STRING,
+        int: DbSqlType.INTEGER,
+        float: DbSqlType.FLOAT,
+        datetime.datetime: DbSqlType.TIMESTAMP,
+        bool: DbSqlType.BOOLEAN,
+    }
+    newParams = copy.deepcopy(params)
+    for param in newParams:
+        if not param.type:
+            if type(param.value) in type_lookup_table:
+                param.type = type_lookup_table[type(param.value)]
+            else:
+                raise ValueError("Parameter type cannot be inferred")
+        param.value = str(param.value)
+    return newParams
+
+
+def named_parameters_to_tsparkparams(parameters: Union[List[Any], Dict[str, str]]):
+    tspark_params = []
+    if isinstance(parameters, dict):
+        dbsql_params = named_parameters_to_dbsqlparams_v1(parameters)
+    else:
+        dbsql_params = named_parameters_to_dbsqlparams_v2(parameters)
+    inferred_type_parameters = infer_types(dbsql_params)
+    for param in inferred_type_parameters:
+        tspark_params.append(
+            TSparkParameter(
+                type=param.type.value,
+                name=param.name,
+                value=TSparkParameterValue(stringValue=param.value),
+            )
+        )
+    return tspark_params

tests/unit/test_parameters.py

@@ -0,0 +1,75 @@
+from databricks.sql.utils import (
+    named_parameters_to_tsparkparams,
+    infer_types,
+    named_parameters_to_dbsqlparams_v1,
+    named_parameters_to_dbsqlparams_v2,
+)
+from databricks.sql.thrift_api.TCLIService.ttypes import (
+    TSparkParameter,
+    TSparkParameterValue,
+)
+from databricks.sql.utils import DbSqlParameter, DbSqlType
+import pytest
+
+
+class TestTSparkParameterConversion(object):
+    def test_conversion_e2e(self):
+        """This behaviour falls back to Python's default string formatting of numbers"""
+        assert named_parameters_to_tsparkparams(
+            ["a", 1, True, 1.0, DbSqlParameter(value="1.0", type=DbSqlType.DECIMAL)]
+        ) == [
+            TSparkParameter(
+                name="", type="STRING", value=TSparkParameterValue(stringValue="a")
+            ),
+            TSparkParameter(
+                name="", type="INTEGER", value=TSparkParameterValue(stringValue="1")
+            ),
+            TSparkParameter(
+                name="", type="BOOLEAN", value=TSparkParameterValue(stringValue="True")
+            ),
+            TSparkParameter(
+                name="", type="FLOAT", value=TSparkParameterValue(stringValue="1.0")
+            ),
+            TSparkParameter(
+                name="", type="DECIMAL", value=TSparkParameterValue(stringValue="1.0")
+            ),
+        ]
+
+    def test_basic_conversions_v1(self):
+        # Test legacy codepath
+        assert named_parameters_to_dbsqlparams_v1({"1": 1, "2": "foo", "3": 2.0}) == [
+            DbSqlParameter("1", 1),
+            DbSqlParameter("2", "foo"),
+            DbSqlParameter("3", 2.0),
+        ]
+
+    def test_basic_conversions_v2(self):
+        # Test interspersing named params with unnamed
+        assert named_parameters_to_dbsqlparams_v2(
+            [DbSqlParameter("1", 1.0, DbSqlType.DECIMAL), 5, DbSqlParameter("3", "foo")]
+        ) == [
+            DbSqlParameter("1", 1.0, DbSqlType.DECIMAL),
+            DbSqlParameter("", 5),
+            DbSqlParameter("3", "foo"),
+        ]
+
+    def test_type_inference(self):
+        with pytest.raises(ValueError):
+            infer_types([DbSqlParameter("", None)])
+        with pytest.raises(ValueError):
+            infer_types([DbSqlParameter("", {1: 1})])
+        assert infer_types([DbSqlParameter("", 1)]) == [
+            DbSqlParameter("", "1", DbSqlType.INTEGER)
+        ]
+        assert infer_types([DbSqlParameter("", True)]) == [
+            DbSqlParameter("", "True", DbSqlType.BOOLEAN)
+        ]
+        assert infer_types([DbSqlParameter("", 1.0)]) == [
+            DbSqlParameter("", "1.0", DbSqlType.FLOAT)
+        ]
+        assert infer_types([DbSqlParameter("", "foo")]) == [
+            DbSqlParameter("", "foo", DbSqlType.STRING)
+        ]
+        assert infer_types([DbSqlParameter("", 1.0, DbSqlType.DECIMAL)]) == [
+            DbSqlParameter("", "1.0", DbSqlType.DECIMAL)
+        ]