From cd6965b8b175fc5d529b95d67a3aa53a788f795a Mon Sep 17 00:00:00 2001
From: Yasuhisa Yoshida <syou6162@gmail.com>
Date: Mon, 4 Nov 2024 21:53:29 +0900
Subject: [PATCH] add case for incremental model

---
 .../macros/materializations/incremental.sql   |  1 +
 tests/functional/test_data_profile_scan.py    | 74 +++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/dbt/include/bigquery/macros/materializations/incremental.sql b/dbt/include/bigquery/macros/materializations/incremental.sql
index 935280d63..b5708545c 100644
--- a/dbt/include/bigquery/macros/materializations/incremental.sql
+++ b/dbt/include/bigquery/macros/materializations/incremental.sql
@@ -170,6 +170,7 @@
   {% do apply_grants(target_relation, grant_config, should_revoke) %}
 
   {% do persist_docs(target_relation, model) %}
+  {% do adapter.create_or_update_data_profile_scan(model) %}
 
   {%- if tmp_relation_exists -%}
     {{ adapter.drop_relation(tmp_relation) }}
diff --git a/tests/functional/test_data_profile_scan.py b/tests/functional/test_data_profile_scan.py
index 80d3332e8..01a4c1839 100644
--- a/tests/functional/test_data_profile_scan.py
+++ b/tests/functional/test_data_profile_scan.py
@@ -52,6 +52,26 @@
         cron: "CRON_TZ=Asia/New_York 0 9 * * *"
 """
 
+INCREMENTAL_MODEL_CONTENT = """
+{{
+    config(
+        materialized="incremental",
+    )
+}}
+
+{% if not is_incremental() %}
+
+    select 10 as id, cast('2020-01-01 01:00:00' as datetime) as date_hour union all
+    select 30 as id, cast('2020-01-01 02:00:00' as datetime) as date_hour
+
+{% else %}
+
+    select 20 as id, cast('2020-01-01 01:00:00' as datetime) as date_hour union all
+    select 40 as id, cast('2020-01-01 02:00:00' as datetime) as date_hour
+
+{% endif %}
+"""
+
 
 class TestDataProfileScanWithProjectProfileScanSetting:
     @pytest.fixture(scope="class")
@@ -362,3 +382,57 @@ def test_create_data_profile_scan(self, project):
                 )
                 labels_to_be_created = list(ORIGINAL_LABELS.keys())
                 assert set(table.labels.keys()) == set(labels_to_be_created)
+
+
+class TestDataProfileScanWithIncrementalModel:
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "models": {
+                "+labels": ORIGINAL_LABELS,
+                "+data_profile_scan": {
+                    "location": SCAN_LOCATION,
+                    "scan_id": SCAN_ID,
+                    "sampling_percent": 10,
+                    "row_filter": "TRUE",
+                },
+            },
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            f"{MODEL_NAME}.sql": INCREMENTAL_MODEL_CONTENT,
+            f"{MODEL_NAME}.yml": YAML_CONTENT,
+        }
+
+    def test_create_data_profile_scan(self, project):
+        with patch("google.cloud.dataplex_v1.DataScanServiceClient") as MockDataScanClient:
+            mock_data_scan_client = MockDataScanClient.return_value
+
+            results = run_dbt()
+            assert len(results) == 1
+
+            mock_data_scan_client.create_data_scan.assert_called_once()
+            mock_data_scan_client.run_data_scan.assert_called_once()
+
+            def list_data_scans_mock(parent):
+                mock_scan = MagicMock()
+                mock_scan.name = SCAN_ID
+                return [mock_scan]
+
+            mock_data_scan_client.list_data_scans.side_effect = list_data_scans_mock
+
+            results = run_dbt()
+            assert len(results) == 1
+            mock_data_scan_client.update_data_scan.assert_called_once()
+
+            relation: BigQueryRelation = relation_from_name(project.adapter, MODEL_NAME)
+            with get_connection(project.adapter) as conn:
+                table = conn.handle.get_table(
+                    project.adapter.connections.get_bq_table(
+                        relation.database, relation.schema, relation.table
+                    )
+                )
+                labels_to_be_created = PROFILE_SCAN_LABELS + list(ORIGINAL_LABELS.keys())
+                assert set(table.labels.keys()) == set(labels_to_be_created)