fix(ingestion/airflow-plugin): updated testcase

datahub-project · Sep 16, 2024 · e5c46b0 · e5c46b0
1 parent 23f73e0
commit e5c46b0
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 8 deletions.
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/operators/datahub.py
@@ -75,9 +75,9 @@ def _render_template_fields(self, field_value, context, jinja_env):
             for item in field_value:
                 self._render_template_fields(item, context, jinja_env)
         elif isinstance(field_value, str):
-            return self.render_template(field_value, context, jinja_env)
+            return super().render_template(field_value, context, jinja_env)
         else:
-            return self.render_template(field_value, context, jinja_env)
+            return super().render_template(field_value, context, jinja_env)
         return field_value
 
     def execute(self, context):

diff --git a/...ules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py b/...ules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py
@@ -50,6 +50,16 @@
                     ],
                 ),
             ),
+            MetadataChangeEvent(
+                proposedSnapshot=DatasetSnapshotClass(
+                    urn="urn:li:dataset:(urn:li:dataPlatform:hive,datahub.example.lineage_example_{{ ts_nodash }},DEV)",
+                    aspects=[
+                        DatasetPropertiesClass(
+                            customProperties={"jinjaTemplate": "{{ ds }}"}
+                        )
+                    ],
+                ),
+            ),
         ],
         datahub_conn_id="datahub_file_default",
     )
diff --git a/...flow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag.json b/...flow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag.json
@@ -242,7 +242,7 @@
             "name": "datahub_emitter_operator_jinja_template_dag_datahub_emitter_operator_jinja_template_dag_task_manual_run_test",
             "type": "BATCH_AD_HOC",
             "created": {
-                "time": 1726471968317,
+                "time": 1726473108543,
                 "actor": "urn:li:corpuser:datahub"
             }
         }
@@ -267,7 +267,7 @@
     "aspectName": "dataProcessInstanceRunEvent",
     "aspect": {
         "json": {
-            "timestampMillis": 1726471968317,
+            "timestampMillis": 1726473108543,
             "partitionSpec": {
                 "partition": "FULL_TABLE_SNAPSHOT",
                 "type": "FULL_TABLE"
@@ -309,6 +309,23 @@
         }
     }
 },
+{
+    "proposedSnapshot": {
+        "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,datahub.example.lineage_example_20230927T213438,DEV)",
+            "aspects": [
+                {
+                    "com.linkedin.pegasus2avro.dataset.DatasetProperties": {
+                        "customProperties": {
+                            "jinjaTemplate": "2023-09-27"
+                        },
+                        "tags": []
+                    }
+                }
+            ]
+        }
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,datahub_emitter_operator_jinja_template_dag,prod),datahub_emitter_operator_jinja_template_dag_task)",
@@ -409,7 +426,7 @@
     "aspectName": "dataProcessInstanceRunEvent",
     "aspect": {
         "json": {
-            "timestampMillis": 1726471968472,
+            "timestampMillis": 1726473108751,
             "partitionSpec": {
                 "partition": "FULL_TABLE_SNAPSHOT",
                 "type": "FULL_TABLE"

diff --git a/...s/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag_no_dag_listener.json b/...s/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag_no_dag_listener.json
@@ -242,7 +242,7 @@
             "name": "datahub_emitter_operator_jinja_template_dag_datahub_emitter_operator_jinja_template_dag_task_manual_run_test",
             "type": "BATCH_AD_HOC",
             "created": {
-                "time": 1726471924787,
+                "time": 1726473074519,
                 "actor": "urn:li:corpuser:datahub"
             }
         }
@@ -267,7 +267,7 @@
     "aspectName": "dataProcessInstanceRunEvent",
     "aspect": {
         "json": {
-            "timestampMillis": 1726471924787,
+            "timestampMillis": 1726473074519,
             "partitionSpec": {
                 "partition": "FULL_TABLE_SNAPSHOT",
                 "type": "FULL_TABLE"
@@ -309,6 +309,23 @@
         }
     }
 },
+{
+    "proposedSnapshot": {
+        "com.linkedin.pegasus2avro.metadata.snapshot.DatasetSnapshot": {
+            "urn": "urn:li:dataset:(urn:li:dataPlatform:hive,datahub.example.lineage_example_20230927T213438,DEV)",
+            "aspects": [
+                {
+                    "com.linkedin.pegasus2avro.dataset.DatasetProperties": {
+                        "customProperties": {
+                            "jinjaTemplate": "2023-09-27"
+                        },
+                        "tags": []
+                    }
+                }
+            ]
+        }
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,datahub_emitter_operator_jinja_template_dag,prod),datahub_emitter_operator_jinja_template_dag_task)",
@@ -409,7 +426,7 @@
     "aspectName": "dataProcessInstanceRunEvent",
     "aspect": {
         "json": {
-            "timestampMillis": 1726471924940,
+            "timestampMillis": 1726473074625,
             "partitionSpec": {
                 "partition": "FULL_TABLE_SNAPSHOT",
                 "type": "FULL_TABLE"