(3/x) Add e2e tests for DECIMAL inferrence

Signed-off-by: Jesse Whitehouse <[email protected]>

tests/e2e/common/parameterized_query_tests.py

@@ -1,13 +1,20 @@
 import datetime
 from decimal import Decimal
+from enum import Enum
 from typing import Dict, List, Tuple, Union
 
 import pytz
 
-from databricks.sql.client import Connection
+from databricks.sql.exc import DatabaseError
 from databricks.sql.utils import DbSqlParameter, DbSqlType
 
 
+class MyCustomDecimalType(Enum):
+    DECIMAL_38_0 = "DECIMAL(38,0)"
+    DECIMAL_38_2 = "DECIMAL(38,2)"
+    DECIMAL_18_9 = "DECIMAL(18,9)"
+
+
 class PySQLParameterizedQueryTestSuiteMixin:
     """Namespace for tests of server-side parameterized queries"""
 
@@ -63,6 +70,11 @@ def test_primitive_inferred_string(self):
         result = self._get_one_result(self.QUERY, params)
         assert result.col == "Hello"
 
+    def test_primitive_inferred_decimal(self):
+        params = {"p": Decimal("1234.56")}
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == Decimal("1234.56")
+
     def test_dbsqlparam_inferred_bool(self):
 
         params = [DbSqlParameter(name="p", value=True, type=None)]
@@ -103,6 +115,11 @@ def test_dbsqlparam_inferred_string(self):
         result = self._get_one_result(self.QUERY, params)
         assert result.col == "Hello"
 
+    def test_dbsqlparam_inferred_decimal(self):
+        params = [DbSqlParameter(name="p", value=Decimal("1234.56"), type=None)]
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == Decimal("1234.56")
+
     def test_dbsqlparam_explicit_bool(self):
 
         params = [DbSqlParameter(name="p", value=True, type=DbSqlType.BOOLEAN)]
@@ -142,3 +159,51 @@ def test_dbsqlparam_explicit_string(self):
         params = [DbSqlParameter(name="p", value="Hello", type=DbSqlType.STRING)]
         result = self._get_one_result(self.QUERY, params)
         assert result.col == "Hello"
+
+    def test_dbsqlparam_explicit_decimal(self):
+        params = [
+            DbSqlParameter(name="p", value=Decimal("1234.56"), type=DbSqlType.DECIMAL)
+        ]
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == Decimal("1234.56")
+
+    def test_dbsqlparam_custom_explicit_decimal_38_0(self):
+
+        # This DECIMAL can be contained in a DECIMAL(38,0) column in Databricks
+        value = Decimal("12345678912345678912345678912345678912")
+        params = [
+            DbSqlParameter(name="p", value=value, type=MyCustomDecimalType.DECIMAL_38_0)
+        ]
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == value
+
+    def test_dbsqlparam_custom_explicit_decimal_38_2(self):
+
+        # This DECIMAL can be contained in a DECIMAL(38,2) column in Databricks
+        value = Decimal("123456789123456789123456789123456789.12")
+        params = [
+            DbSqlParameter(name="p", value=value, type=MyCustomDecimalType.DECIMAL_38_2)
+        ]
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == value
+
+    def test_dbsqlparam_custom_explicit_decimal_18_9(self):
+
+        # This DECIMAL can be contained in a DECIMAL(38,2) column in Databricks
+        value = Decimal("123456789.123456789")
+        params = [
+            DbSqlParameter(name="p", value=value, type=MyCustomDecimalType.DECIMAL_18_9)
+        ]
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == value
+
+    def test_primitive_inferred_decimal_raises_exception(self):
+        """The default precision is DECIMAL(6,2). Without a custom DbsqlParameter type, the value will be rounded"""
+
+        # This DECIMAL would require DECIMAL(10,2) but the default is DECIMAL(6,2)
+        params = {"p": Decimal("12345678.91")}
+
+        with self.assertRaises(
+            DatabaseError, msg="cannot be cast to DECIMAL(6,2) because it is malformed"
+        ):
+            result = self._get_one_result(self.QUERY, params)