From e5c46b043e88bfef7e0cbd8327a4b5dfbe160108 Mon Sep 17 00:00:00 2001 From: Dushyant Bhalgami Date: Mon, 16 Sep 2024 09:55:40 +0200 Subject: [PATCH] fix(ingestion/airflow-plugin): updated testcase --- .../operators/datahub.py | 4 ++-- ...hub_emitter_operator_jinja_template_dag.py | 10 ++++++++ ...b_emitter_operator_jinja_template_dag.json | 23 ++++++++++++++++--- ...or_jinja_template_dag_no_dag_listener.json | 23 ++++++++++++++++--- 4 files changed, 52 insertions(+), 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 index 8c79eb7c56b4d3..f5c34da999babc 100644 --- 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/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py index 1c9daa981eb5e0..79fff504efc8d3 100644 --- a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/datahub_emitter_operator_jinja_template_dag.py +++ b/metadata-ingestion-modules/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/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag.json index 7684e24f11c750..aa97ed10652e23 100644 --- a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag.json +++ b/metadata-ingestion-modules/airflow-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/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag_no_dag_listener.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag_no_dag_listener.json index dee46468c202be..a163acbefdbc03 100644 --- a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_datahub_emitter_operator_jinja_template_dag_no_dag_listener.json +++ b/metadata-ingestion-modules/airflow-plugin/tests/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"