Add tests and validation

getsentry · May 12, 2024 · c08623f · c08623f
1 parent b813c29
commit c08623f
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 2 deletions.
diff --git a/snuba/datasets/configuration/json_schema.py b/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", "DateTime64", "UUID", "IPv4", "IPv6"]},
     args={
         "type": "object",
         "properties": {},

diff --git a/snuba/datasets/configuration/utils.py b/snuba/datasets/configuration/utils.py
@@ -6,6 +6,7 @@
     Array,
     Column,
     DateTime,
+    DateTime64,
     Enum,
     Float,
     Nested,
@@ -66,6 +67,7 @@ def get_mandatory_condition_checkers(
     **NUMBER_COLUMN_TYPES,
     "String": String,
     "DateTime": DateTime,
+    "DateTime64": DateTime64,
     "UUID": UUID,
     "IPv4": IPv4,
     "IPv6": IPv6,

diff --git a/snuba/migrations/parse_schema.py b/snuba/migrations/parse_schema.py
@@ -14,6 +14,7 @@
     ColumnType,
     Date,
     DateTime,
+    DateTime64,
     Enum,
     FixedString,
     Float,
@@ -29,7 +30,7 @@
     type             = primitive / lowcardinality / agg / nullable / array
     primitive        = basic_type / uint / float / fixedstring / enum
     # DateTime must come before Date
-    basic_type       = "DateTime" / "Date" / "IPv4" / "IPv6" / "String" / "UUID"
+    basic_type       = "DateTime" / "DateTime64" / "Date" / "IPv4" / "IPv6" / "String" / "UUID"
     uint             = "UInt" uint_size
     uint_size        = "8" / "16" / "32" / "64"
     float            = "Float" float_size
@@ -71,6 +72,7 @@ def merge_modifiers(
 _TYPES: dict[str, type[ColumnType[MigrationModifiers]]] = {
     "Date": Date,
     "DateTime": DateTime,
+    "DateTime64": DateTime64,
     "IPv4": IPv4,
     "IPv6": IPv6,
     "String": String,

diff --git a/tests/clickhouse/test_columns.py b/tests/clickhouse/test_columns.py
@@ -11,6 +11,7 @@
     ColumnType,
     Date,
     DateTime,
+    DateTime64,
     Enum,
     FixedString,
     Float,
@@ -86,6 +87,13 @@
         "DateTime",
         id="datetimes",
     ),
+    pytest.param(
+        DateTime64(),
+        DateTime64(),
+        DateTime64(Modifier(nullable=True)),
+        "DateTime64",
+        id="datetime64s",
+    ),
     pytest.param(
         Array(String(Modifier(nullable=True))),
         Array(String()),

diff --git a/tests/datasets/configuration/test_storage_loader.py b/tests/datasets/configuration/test_storage_loader.py
@@ -8,6 +8,7 @@
     Array,
     Column,
     DateTime,
+    DateTime64,
     Enum,
     Float,
     Nested,
@@ -120,6 +121,7 @@ def test_column_parser(self) -> None:
             {"name": "float_col", "type": "Float", "args": {"size": 32}},
             {"name": "string_col", "type": "String"},
             {"name": "time_col", "type": "DateTime"},
+            {"name": "time64_col", "type": "DateTime64"},
             {
                 "name": "nested_col",
                 "type": "Nested",
@@ -171,6 +173,7 @@ def test_column_parser(self) -> None:
             Column("float_col", Float(32)),
             Column("string_col", String()),
             Column("time_col", DateTime()),
+            Column("time64_col", DateTime64()),
             Column("nested_col", Nested([Column("sub_col", UInt(64))])),
             Column("func_col", AggregateFunction("uniqCombined64", [UInt(64)])),
             Column(