From c08623ff703165e084c8b81eb7c6f788a4cb134f Mon Sep 17 00:00:00 2001
From: Pierre Massat <pierre.massat@sentry.io>
Date: Sun, 12 May 2024 11:27:57 -0400
Subject: [PATCH] Add tests and validation

---
 snuba/datasets/configuration/json_schema.py         | 2 +-
 snuba/datasets/configuration/utils.py               | 2 ++
 snuba/migrations/parse_schema.py                    | 4 +++-
 tests/clickhouse/test_columns.py                    | 8 ++++++++
 tests/datasets/configuration/test_storage_loader.py | 3 +++
 5 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/snuba/datasets/configuration/json_schema.py b/snuba/datasets/configuration/json_schema.py
index 573ebef57e..68702fb715 100644
--- 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
index 725390293c..a647bce23a 100644
--- 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
index 82db092d09..7cc556a72f 100644
--- 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
index b01508bdd0..f1fc697f56 100644
--- 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
index a728cc3df4..e0ed14fc2f 100644
--- 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(