From e5166ee2b6ffb4db71e1418e1f7e02dfa222f66f Mon Sep 17 00:00:00 2001
From: Connor Charles <Connor.Charles@autotrader.co.uk>
Date: Sun, 9 Oct 2022 20:15:56 +0100
Subject: [PATCH] Support list of strings for incremental unique_key

---
 dbt_dry_run/models/manifest.py                |  2 +-
 dbt_dry_run/node_runner/snapshot_runner.py    |  5 +++
 .../test/node_runner/test_snapshot_runner.py  | 40 +++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/dbt_dry_run/models/manifest.py b/dbt_dry_run/models/manifest.py
index 04e3cd9..686f83d 100644
--- a/dbt_dry_run/models/manifest.py
+++ b/dbt_dry_run/models/manifest.py
@@ -40,7 +40,7 @@ class NodeConfig(BaseModel):
     materialized: str
     on_schema_change: Optional[OnSchemaChange]
     sql_header: Optional[str]
-    unique_key: Optional[str]
+    unique_key: Optional[Union[str, List[str]]]
     updated_at: Optional[str]
     strategy: Union[None, Literal["timestamp", "check"]]
     check_cols: Optional[Union[Literal["all"], List[str]]]
diff --git a/dbt_dry_run/node_runner/snapshot_runner.py b/dbt_dry_run/node_runner/snapshot_runner.py
index 74f317a..1943b25 100644
--- a/dbt_dry_run/node_runner/snapshot_runner.py
+++ b/dbt_dry_run/node_runner/snapshot_runner.py
@@ -45,6 +45,11 @@ class SnapshotRunner(NodeRunner):
     def _validate_snapshot_config(node: Node, result: DryRunResult) -> DryRunResult:
         if not result.table:
             raise ValueError("Can't validate result without table")
+        if isinstance(node.config.unique_key, list):
+            raise RuntimeError(
+                f"Cannot dry run node '{node.unique_id}' because it is a snapshot"
+                f" with a list of unique keys '{node.config.unique_key}'"
+            )
         if node.config.unique_key not in result.table.field_names:
             exception = SnapshotConfigException(
                 f"Missing `unique_key` column '{node.config.unique_key}'"
diff --git a/dbt_dry_run/test/node_runner/test_snapshot_runner.py b/dbt_dry_run/test/node_runner/test_snapshot_runner.py
index 2d0ef3d..de4efd3 100644
--- a/dbt_dry_run/test/node_runner/test_snapshot_runner.py
+++ b/dbt_dry_run/test/node_runner/test_snapshot_runner.py
@@ -1,5 +1,7 @@
 from unittest.mock import MagicMock
 
+import pytest
+
 from dbt_dry_run.literals import enable_test_example_values
 from dbt_dry_run.models import BigQueryFieldType, Table, TableField
 from dbt_dry_run.models.manifest import NodeConfig
@@ -240,3 +242,41 @@ def test_snapshot_with_timestamp_strategy_with_missing_updated_at_column() -> No
     result = model_runner.run(node)
     mock_sql_runner.query.assert_called_with(node.compiled_sql)
     assert result.status == DryRunStatus.FAILURE
+
+
+def test_snapshot_with_list_of_unique_key_columns_raises_error() -> None:
+    """
+    This isn't currently supported by dbt-core but this could change given it was added for incremental here:
+    https://github.com/dbt-labs/dbt-core/pull/4618
+    """
+    mock_sql_runner = MagicMock()
+    expected_table = Table(
+        fields=[
+            TableField(
+                name="a",
+                type=BigQueryFieldType.STRING,
+            ),
+            TableField(name="last_updated_col", type=BigQueryFieldType.TIMESTAMP),
+        ]
+    )
+    mock_sql_runner.query.return_value = (DryRunStatus.SUCCESS, expected_table, None)
+
+    node = SimpleNode(
+        unique_id="node1",
+        depends_on=[],
+        resource_type=ManifestScheduler.SNAPSHOT,
+        table_config=NodeConfig(
+            unique_key=["a", "b"],
+            strategy="timestamp",
+            materialized="snapshot",
+            updated_at="wrong_last_updated_col",
+        ),
+    ).to_node()
+    node.depends_on.deep_nodes = []
+
+    results = Results()
+
+    model_runner = SnapshotRunner(mock_sql_runner, results)
+
+    with pytest.raises(RuntimeError, match="Cannot dry run node"):
+        model_runner.run(node)