feat: Add a Bool column type (#6861)

Author: phacops

snuba/clickhouse/columns.py

@@ -2,7 +2,7 @@
 
 from typing import Sequence, Union
 
-from snuba.utils.schemas import UUID, AggregateFunction, Any, Array, Column
+from snuba.utils.schemas import UUID, AggregateFunction, Any, Array, Bool, Column
 from snuba.utils.schemas import ColumnSet as BaseColumnSet
 from snuba.utils.schemas import (
     ColumnType,
@@ -35,6 +35,7 @@
     "Any",
     "AggregateFunction",
     "Array",
+    "Bool",
     "Column",
     "ColumnSet",
     "ColumnType",

snuba/datasets/configuration/json_schema.py

@@ -141,7 +141,7 @@ def del_name_field(column_schema: dict[str, Any]) -> dict[str, Any]:
 
 
 NO_ARG_SCHEMA = make_column_schema(
-    column_type={"enum": ["String", "DateTime", "UUID", "IPv4", "IPv6"]},
+    column_type={"enum": ["String", "DateTime", "UUID", "IPv4", "IPv6", "Bool"]},
     args={
         "type": "object",
         "properties": {},

snuba/datasets/configuration/utils.py

@@ -20,6 +20,7 @@
 from snuba.utils.schemas import (
     UUID,
     AggregateFunction,
+    Bool,
     ColumnType,
     FixedString,
     Int,
@@ -75,6 +76,7 @@ def get_mandatory_condition_checkers(
     "UUID": UUID,
     "IPv4": IPv4,
     "IPv6": IPv6,
+    "Bool": Bool,
 }
 
 

snuba/migrations/parse_schema.py

@@ -11,6 +11,7 @@
     UUID,
     AggregateFunction,
     Array,
+    Bool,
     ColumnType,
     Date,
     DateTime,
@@ -35,7 +36,7 @@
     # datetime64 needs to be before basic_type to not be parsed as DateTime
     primitive        = datetime64 / basic_type / uint / int / float / fixedstring / enum
     # DateTime must come before Date
-    basic_type       = "DateTime" / "Date" / "IPv4" / "IPv6" / "String" / "UUID"
+    basic_type       = "DateTime" / "Date" / "IPv4" / "IPv6" / "String" / "UUID" / "Bool"
     uint             = "UInt" uint_size
     int              = "Int" uint_size
     uint_size        = "8" / "16" / "32" / "64" / "128"
@@ -82,6 +83,7 @@ def merge_modifiers(
 
 
 _TYPES: dict[str, type[ColumnType[MigrationModifiers]]] = {
+    "Bool": Bool,
     "Date": Date,
     "DateTime": DateTime,
     "IPv4": IPv4,

snuba/utils/schemas.py

@@ -852,3 +852,23 @@ def _valid_tuple(self, tuple_column: Tuple[AnyType], value: tuple[Any]) -> bool:
                 f"Invalid value {el} for column type {tuple_column.types[i]}"
             )
         return True
+
+
+class Bool(ColumnType[TModifiers]):
+    def __init__(self, modifiers: Optional[TModifiers] = None) -> None:
+        super().__init__(modifiers)
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            self.__class__ == other.__class__
+            and self.get_modifiers() == cast(Int[TModifiers], other).get_modifiers()
+        )
+
+    def _for_schema_impl(self) -> str:
+        return "Bool"
+
+    def set_modifiers(self, modifiers: Optional[TModifiers]) -> Bool[TModifiers]:
+        return Bool(modifiers=modifiers)
+
+    def get_raw(self) -> Bool[TModifiers]:
+        return Bool()

tests/clickhouse/test_columns.py

@@ -7,6 +7,7 @@
     UUID,
     AggregateFunction,
     Array,
+    Bool,
     Column,
     ColumnType,
     Date,
@@ -173,6 +174,13 @@
         "Nullable(Enum('a' = 1, 'b' = 2))",
         id="enums",
     ),
+    pytest.param(
+        Bool(Modifier(nullable=True)),
+        Bool(),
+        Bool(),
+        "Nullable(Bool)",
+        id="bools",
+    ),
 ]
 
 

tests/datasets/configuration/test_storage_loader.py

@@ -6,6 +6,7 @@
 
 from snuba.clickhouse.columns import (
     Array,
+    Bool,
     Column,
     DateTime,
     DateTime64,
@@ -153,6 +154,7 @@ def test_column_parser(self) -> None:
             {"name": "int_col", "type": "UInt", "args": {"size": 64}},
             {"name": "float_col", "type": "Float", "args": {"size": 32}},
             {"name": "string_col", "type": "String"},
+            {"name": "bool_col", "type": "Bool"},
             {"name": "time_col", "type": "DateTime"},
             {
                 "name": "time64_col",
@@ -209,6 +211,7 @@ def test_column_parser(self) -> None:
             Column("int_col", UInt(64)),
             Column("float_col", Float(32)),
             Column("string_col", String()),
+            Column("bool_col", Bool()),
             Column("time_col", DateTime()),
             Column("time64_col", DateTime64(3, "America/New_York")),
             Column("nested_col", Nested([Column("sub_col", UInt(64))])),

tests/migrations/test_parse_schema.py

@@ -6,6 +6,7 @@
     UUID,
     AggregateFunction,
     Array,
+    Bool,
     ColumnType,
     Date,
     DateTime,
@@ -38,6 +39,7 @@
     (("UInt32", "", "", ""), UInt(32)),
     (("UInt128", "", "", ""), UInt(128)),
     (("UUID", "", "", ""), UUID()),
+    (("Bool", "", "", ""), Bool()),
     # Aggregate functions
     (
         ("AggregateFunction(uniq, UInt8)", "", "", ""),