feat: Fir 16854 add boolean support (#231)

Author: stepansergeevitch

src/firebolt/async_db/_types.py

@@ -200,6 +200,9 @@ class _InternalType(Enum):
     TimestampNtz = "TimestampNtz"
     TimestampTz = "TimestampTz"
 
+    # BOOLEAN
+    Boolean = "boolean"
+
     # Nullable(Nothing)
     Nothing = "Nothing"
 
@@ -224,6 +227,7 @@ def python_type(self) -> type:
             _InternalType.DateTime: datetime,
             _InternalType.TimestampNtz: datetime,
             _InternalType.TimestampTz: datetime,
+            _InternalType.Boolean: bool,
             # For simplicity, this could happen only during 'select null' query
             _InternalType.Nothing: str,
         }
@@ -285,6 +289,10 @@ def parse_value(
         if not isinstance(value, str):
             raise DataError(f"Invalid datetime value {value}: str expected")
         return parse_datetime(value)
+    if ctype is bool:
+        if not isinstance(value, (bool, int)):
+            raise DataError(f"Invalid boolean value {value}: bool or int expected")
+        return bool(value)
     if isinstance(ctype, DECIMAL):
         assert isinstance(value, (str, int))
         return Decimal(value)
@@ -304,7 +312,7 @@ def parse_value(
 def format_value(value: ParameterType) -> str:
     """For Python value to be used in a SQL query."""
     if isinstance(value, bool):
-        return str(int(value))
+        return "true" if value else "false"
     if isinstance(value, (int, float, Decimal)):
         return str(value)
     elif isinstance(value, str):

tests/integration/dbapi/async/test_queries_async.py

@@ -138,9 +138,11 @@ async def test_select(
         assert (
             await c.execute(f"SET advanced_mode=1") == -1
         ), "Invalid set statment row count"
+        # For TimestampTz test
         assert (
             await c.execute(f"SET time_zone={timezone_name}") == -1
         ), "Invalid set statment row count"
+
         assert await c.execute(all_types_query) == 1, "Invalid row count returned"
         assert c.rowcount == 1, "Invalid rowcount value"
         data = await c.fetchall()
@@ -166,6 +168,37 @@ async def test_select(
         )
 
 
+async def test_boolean(
+    connection: Connection,
+) -> None:
+    """Select handles boolean properly."""
+    with connection.cursor() as c:
+        assert (
+            await c.execute(f"SET advanced_mode=1") == -1
+        ), "Invalid set statment row count"
+        # Unfortunately our parser doesn't support string set parameters
+        c._set_parameters.update(
+            {
+                "bool_output_format": "postgres",
+                "output_format_firebolt_type_names": "true",
+            }
+        )
+
+        assert (
+            await c.execute('select true as "bool"') == 1
+        ), "Invalid row count returned"
+        assert c.rowcount == 1, "Invalid rowcount value"
+        data = await c.fetchall()
+        assert len(data) == c.rowcount, "Invalid data length"
+        assert_deep_eq(data, [[True]], "Invalid data")
+        assert (
+            c.description == [Column("bool", bool, None, None, None, None, None)],
+            "Invalid description value",
+        )
+        assert len(data[0]) == len(c.description), "Invalid description length"
+        assert len(await c.fetchall()) == 0, "Redundant data returned by fetchall"
+
+
 @mark.skip("Don't have a good way to test this anymore. FIR-16038")
 @mark.timeout(timeout=400)
 async def test_long_query(

tests/integration/dbapi/conftest.py

@@ -72,7 +72,6 @@ def all_types_query() -> str:
         "CAST('2019-07-31 01:01:01.1234' AS TIMESTAMP_EXT(4)) as \"datetime64\", "
         "CAST('1111-01-05 17:04:42.123456' as timestampntz) as timestampntz, "
         "'1111-01-05 17:04:42.123456'::timestamptz as timestamptz,"
-        'true as "bool",'
         "[1,2,3,4] as \"array\", cast('1231232.123459999990457054844258706536' as "
         'decimal(38,30)) as "decimal", '
         "cast(null as  int) as nullable"
@@ -100,7 +99,6 @@ def all_types_query_description() -> List[Column]:
         Column("datetime64", DATETIME64(4), None, None, None, None, None),
         Column("timestampntz", datetime, None, None, None, None, None),
         Column("timestamptz", datetime, None, None, None, None, None),
-        Column("bool", int, None, None, None, None, None),
         Column("array", ARRAY(int), None, None, None, None, None),
         Column("decimal", DECIMAL(38, 30), None, None, None, None, None),
         Column("nullable", int, None, None, None, None, None),
@@ -138,7 +136,6 @@ def all_types_query_response(timezone_offset_seconds: int) -> List[ColType]:
                 123456,
                 tzinfo=timezone(timedelta(seconds=timezone_offset_seconds)),
             ),
-            1,
             [1, 2, 3, 4],
             Decimal("1231232.123459999990457054844258706536"),
             None,

tests/integration/dbapi/sync/test_queries.py

@@ -122,6 +122,32 @@ def test_select(
         )
 
 
+def test_boolean(
+    connection: Connection,
+) -> None:
+    """Select handles boolean properly."""
+    with connection.cursor() as c:
+        assert c.execute(f"SET advanced_mode=1") == -1, "Invalid set statment row count"
+        # Unfortunately our parser doesn't support string set parameters
+        c._set_parameters.update(
+            {
+                "bool_output_format": "postgres",
+                "output_format_firebolt_type_names": "true",
+            }
+        )
+
+        assert c.execute('select true as "bool"') == 1, "Invalid row count returned"
+        assert c.rowcount == 1, "Invalid rowcount value"
+        data = c.fetchall()
+        assert len(data) == c.rowcount, "Invalid data length"
+        assert_deep_eq(data, [[True]], "Invalid data")
+        assert c.description == [
+            Column("bool", bool, None, None, None, None, None)
+        ], "Invalid description value"
+        assert len(data[0]) == len(c.description), "Invalid description length"
+        assert len(c.fetchall()) == 0, "Redundant data returned by fetchall"
+
+
 @mark.skip("Don't have a good way to test this anymore. FIR-16038")
 @mark.timeout(timeout=400)
 def test_long_query(

tests/unit/async_db/conftest.py

@@ -60,6 +60,7 @@ def types_map() -> Dict[str, type]:
         "Decimal(38,0)": DECIMAL(38, 0),
         # Invalid decimal format
         "Decimal(38)": str,
+        "boolean": bool,
         "SomeRandomNotExistingType": str,
     }
     array_types = {f"Array({k})": ARRAY(v) for k, v in base_types.items()}

tests/unit/async_db/test_typing_format.py

@@ -29,8 +29,8 @@
         (1.123, "1.123"),
         (Decimal("1.123"), "1.123"),
         (Decimal(1.123), "1.1229999999999999982236431605997495353221893310546875"),
-        (True, "1"),
-        (False, "0"),
+        (True, "true"),
+        (False, "false"),
         # Date, datetime
         (date(2022, 1, 10), "'2022-01-10'"),
         (datetime(2022, 1, 10, 1, 1, 1), "'2022-01-10 01:01:01'"),

tests/unit/async_db/test_typing_parse.py

@@ -236,3 +236,15 @@ def test_helpers() -> None:
 
     with raises(NotSupportedError):
         TimeFromTicks(0)
+
+
+def test_parse_value_bool() -> None:
+    """parse_value parses all int values correctly."""
+    assert parse_value(True, bool) == True, "Error parsing boolean: provided true"
+    assert parse_value(False, bool) == False, "Error parsing boolean: provided false"
+    assert parse_value(2, bool) == True, "Error parsing boolean: provided 2"
+    assert parse_value(0, bool) == False, "Error parsing boolean: provided 0"
+    assert parse_value(None, bool) is None, "Error parsing boolean: provided None"
+
+    with raises(DataError):
+        parse_value("true", bool)