Merge pull request #203 from preset-io/pass_certification

feat: allow passing a certification payload

src/preset_cli/cli/superset/sync/dbt/datasets.py

@@ -6,7 +6,7 @@
 
 import json
 import logging
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from sqlalchemy.engine import create_engine
 from sqlalchemy.engine.url import URL as SQLAlchemyURL
@@ -73,6 +73,7 @@ def sync_datasets(  # pylint: disable=too-many-locals, too-many-branches, too-ma
     database: Any,
     disallow_edits: bool,
     external_url_prefix: str,
+    certification: Optional[Dict[str, Any]] = None,
 ) -> List[Any]:
     """
     Read the dbt manifest and import models as datasets with metrics.
@@ -105,9 +106,8 @@ def sync_datasets(  # pylint: disable=too-many-locals, too-many-branches, too-ma
         extra = {
             "unique_id": model["unique_id"],
             "depends_on": "ref('{name}')".format(**model),
-            "certification": {
-                "details": "This table is produced by dbt",
-            },
+            "certification": certification
+            or {"details": "This table is produced by dbt"},
         }
 
         dataset_metrics = []

tests/cli/superset/sync/dbt/datasets_test.py

@@ -275,6 +275,65 @@ def test_sync_datasets_no_metrics(mocker: MockerFixture) -> None:
     )
 
 
+def test_sync_datasets_custom_certification(mocker: MockerFixture) -> None:
+    """
+    Test ``sync_datasets`` with a custom certification.
+    """
+    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,
+        models=models,
+        metrics=[],
+        database={"id": 1, "sqlalchemy_uri": "postgresql://user@host/examples_dev"},
+        disallow_edits=False,
+        external_url_prefix="",
+        certification={"details": "This dataset is synced from dbt Cloud"},
+    )
+    client.create_dataset.assert_has_calls(
+        [
+            mock.call(database=1, schema="public", table_name="messages_channels"),
+        ],
+    )
+    client.update_dataset.assert_has_calls(
+        [
+            mock.call(
+                1,
+                override_columns=True,
+                description="",
+                extra=json.dumps(
+                    {
+                        "unique_id": "model.superset_examples.messages_channels",
+                        "depends_on": "ref('messages_channels')",
+                        "certification": {
+                            "details": "This dataset is synced from dbt Cloud",
+                        },
+                    },
+                ),
+                is_managed_externally=False,
+                metrics=[],
+            ),
+            mock.call(
+                1,
+                override_columns=True,
+                columns=[
+                    {
+                        "column_name": "id",
+                        "description": "Primary key",
+                        "is_dttm": False,
+                        "verbose_name": "id",
+                    },
+                ],
+            ),
+        ],
+    )
+
+
 def test_sync_datasets_new_bq_error(mocker: MockerFixture) -> None:
     """
     Test ``sync_datasets`` when one of the sources is in a different BQ project.