From 3cebf1b3d729486d3bf4e2d55f1c24e0457c6c9c Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Wed, 30 Nov 2022 14:24:45 -0800
Subject: [PATCH] fix: preserve is_dttm on column update

---
 .../cli/superset/sync/dbt/datasets.py         | 25 +++++++++++++------
 tests/cli/superset/sync/dbt/datasets_test.py  | 23 +++++++++++++++++
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/src/preset_cli/cli/superset/sync/dbt/datasets.py b/src/preset_cli/cli/superset/sync/dbt/datasets.py
index d9485871..b0aed55e 100644
--- a/src/preset_cli/cli/superset/sync/dbt/datasets.py
+++ b/src/preset_cli/cli/superset/sync/dbt/datasets.py
@@ -149,14 +149,23 @@ def sync_datasets(  # pylint: disable=too-many-locals, too-many-branches, too-ma
             client.update_dataset(dataset["id"], override_columns=False, **update)
 
         # update column descriptions
-        update = {
-            "columns": [
-                {"column_name": name, "description": column.get("description", "")}
-                for name, column in model.get("columns", {}).items()
-            ],
-        }
-        if update["columns"]:
-            client.update_dataset(dataset["id"], override_columns=True, **update)
+        if columns := model.get("columns"):
+            current_columns = client.get_dataset(dataset["id"])["columns"]
+            for column in current_columns:
+                name = column["column_name"]
+                if name in columns:
+                    column["description"] = columns[name].get("description", "")
+
+                # remove columns that are not part of the update payload
+                for key in ("changed_on", "created_on", "type_generic"):
+                    if key in column:
+                        del column[key]
+
+            client.update_dataset(
+                dataset["id"],
+                override_columns=True,
+                columns=current_columns,
+            )
 
         datasets.append(dataset)
 
diff --git a/tests/cli/superset/sync/dbt/datasets_test.py b/tests/cli/superset/sync/dbt/datasets_test.py
index 44d471b7..9212ec78 100644
--- a/tests/cli/superset/sync/dbt/datasets_test.py
+++ b/tests/cli/superset/sync/dbt/datasets_test.py
@@ -58,6 +58,12 @@ def test_sync_datasets_new(mocker: MockerFixture) -> None:
     client = mocker.MagicMock()
     client.get_datasets.return_value = []
     client.create_dataset.side_effect = [{"id": 1}, {"id": 2}, {"id": 3}]
+    client.get_dataset.return_value = {
+        "columns": [
+            {"column_name": "id", "is_dttm": False, "type_generic": "INTEGER"},
+            {"column_name": "ts", "is_dttm": True, "type_generic": "TIMESTAMP"},
+        ],
+    }
 
     sync_datasets(
         client=client,
@@ -109,6 +115,11 @@ def test_sync_datasets_new(mocker: MockerFixture) -> None:
                     {
                         "column_name": "id",
                         "description": "Primary key",
+                        "is_dttm": False,
+                    },
+                    {
+                        "column_name": "ts",
+                        "is_dttm": True,
                     },
                 ],
             ),
@@ -123,6 +134,9 @@ def test_sync_datasets_no_metrics(mocker: MockerFixture) -> None:
     client = mocker.MagicMock()
     client.get_datasets.return_value = []
     client.create_dataset.side_effect = [{"id": 1}, {"id": 2}, {"id": 3}]
+    client.get_dataset.return_value = {
+        "columns": [{"column_name": "id", "is_dttm": False}],
+    }
 
     sync_datasets(
         client=client,
@@ -160,6 +174,7 @@ def test_sync_datasets_no_metrics(mocker: MockerFixture) -> None:
                     {
                         "column_name": "id",
                         "description": "Primary key",
+                        "is_dttm": False,
                     },
                 ],
             ),
@@ -201,6 +216,9 @@ def test_sync_datasets_existing(mocker: MockerFixture) -> None:
     """
     client = mocker.MagicMock()
     client.get_datasets.side_effect = [[{"id": 1}], [{"id": 2}], [{"id": 3}]]
+    client.get_dataset.return_value = {
+        "columns": [{"column_name": "id", "is_dttm": False}],
+    }
 
     sync_datasets(
         client=client,
@@ -248,6 +266,7 @@ def test_sync_datasets_existing(mocker: MockerFixture) -> None:
                     {
                         "column_name": "id",
                         "description": "Primary key",
+                        "is_dttm": False,
                     },
                 ],
             ),
@@ -280,6 +299,9 @@ def test_sync_datasets_external_url(mocker: MockerFixture) -> None:
     """
     client = mocker.MagicMock()
     client.get_datasets.side_effect = [[{"id": 1}], [{"id": 2}], [{"id": 3}]]
+    client.get_dataset.return_value = {
+        "columns": [{"column_name": "id", "is_dttm": False}],
+    }
 
     sync_datasets(
         client=client,
@@ -331,6 +353,7 @@ def test_sync_datasets_external_url(mocker: MockerFixture) -> None:
                     {
                         "column_name": "id",
                         "description": "Primary key",
+                        "is_dttm": False,
                     },
                 ],
             ),