Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix(ingestion/airflow-plugin): pipeline tasks discoverable in search #10819

Merged
merged 3 commits into from
Jul 2, 2024
Merged

fix(ingestion/airflow-plugin): pipeline tasks discoverable in search #10819

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bd8c55a
fix(ingestion/airflow-plugin): pipeline tasks discoverable in search
dushayntAW Jul 1, 2024
d746433
fix(ingestion/airflow-plugin): pipeline tasks discoverable in search
dushayntAW Jul 1, 2024
a20058a
fix(ingestion/airflow-plugin): fixed failing CI
dushayntAW Jul 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -536,14 +536,27 @@ def on_dag_start(self, dag_run: "DagRun") -> None:
)
dataflow.emit(self.emitter, callback=self._make_emit_callback())

event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=str(dataflow.urn), aspect=StatusClass(removed=False)
)
Comment on lines +539 to +541
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ensure error handling for MetadataChangeProposalWrapper creation.

Consider adding error handling when creating MetadataChangeProposalWrapper events to ensure that any issues during the creation process are logged and handled gracefully.

+        try:
            event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
                entityUrn=str(dataflow.urn), aspect=StatusClass(removed=False)
            )
+        except Exception as e:
+            logger.error(f"Error creating MetadataChangeProposalWrapper for dataflow: {dataflow.urn}", exc_info=e)
+            return
Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=str(dataflow.urn), aspect=StatusClass(removed=False)
)
try:
event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=str(dataflow.urn), aspect=StatusClass(removed=False)
)
except Exception as e:
logger.error(f"Error creating MetadataChangeProposalWrapper for dataflow: {dataflow.urn}", exc_info=e)
return

self.emitter.emit(event)

for task in dag.tasks:
task_urn = builder.make_data_job_urn_with_flow(
str(dataflow.urn), task.task_id
)
event = MetadataChangeProposalWrapper(
entityUrn=task_urn, aspect=StatusClass(removed=False)
)
self.emitter.emit(event)
Comment on lines +548 to +551
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ensure error handling for task MetadataChangeProposalWrapper creation.

Consider adding error handling when creating MetadataChangeProposalWrapper events for tasks to ensure that any issues during the creation process are logged and handled gracefully.

+            try:
                event = MetadataChangeProposalWrapper(
                    entityUrn=task_urn, aspect=StatusClass(removed=False)
                )
+            except Exception as e:
+                logger.error(f"Error creating MetadataChangeProposalWrapper for task: {task_urn}", exc_info=e)
+                continue

Committable suggestion was skipped due to low confidence.


# emit tags
for tag in dataflow.tags:
tag_urn = builder.make_tag_urn(tag)

event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
event = MetadataChangeProposalWrapper(
entityUrn=tag_urn, aspect=StatusClass(removed=False)
)

self.emitter.emit(event)
Comment on lines +557 to 560
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ensure error handling for tag MetadataChangeProposalWrapper creation.

Consider adding error handling when creating MetadataChangeProposalWrapper events for tags to ensure that any issues during the creation process are logged and handled gracefully.

+            try:
                event = MetadataChangeProposalWrapper(
                    entityUrn=tag_urn, aspect=StatusClass(removed=False)
+            except Exception as e:
+                logger.error(f"Error creating MetadataChangeProposalWrapper for tag: {tag_urn}", exc_info=e)
+                continue

Committable suggestion was skipped due to low confidence.


browse_path_v2_event: MetadataChangeProposalWrapper = (
Expand Down
22 changes: 22 additions & 0 deletions metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,28 @@
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)",
Expand Down
22 changes: 22 additions & 0 deletions ...ion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets_no_dag_listener.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,17 @@
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)",
Expand All @@ -72,6 +83,17 @@
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)",
Expand Down
33 changes: 33 additions & 0 deletions metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,28 @@
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)",
Expand Down Expand Up @@ -102,6 +124,17 @@
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)",
Expand Down
33 changes: 33 additions & 0 deletions ...stion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag_no_dag_listener.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,17 @@
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)",
Expand All @@ -73,6 +84,17 @@
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)",
Expand Down Expand Up @@ -102,6 +124,17 @@
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)",
Expand Down
22 changes: 22 additions & 0 deletions ...ata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_snowflake_operator.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,28 @@
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,snowflake_operator,prod)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,snowflake_operator,prod),transform_cost_table)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,snowflake_operator,prod)",
Expand Down
66 changes: 66 additions & 0 deletions metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,28 @@
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataFlow",
"entityUrn": "urn:li:dataFlow:(airflow,sqlite_operator,prod)",
Expand Down Expand Up @@ -102,6 +124,39 @@
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),transform_cost_table)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_processed_costs)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),populate_cost_table)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),create_cost_table)",
Expand Down Expand Up @@ -155,6 +210,17 @@
}
}
},
{
"entityType": "dataJob",
"entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,sqlite_operator,prod),cleanup_costs)",
"changeType": "UPSERT",
"aspectName": "status",
"aspect": {
"json": {
"removed": false
}
}
},
{
"entityType": "dataset",
"entityUrn": "urn:li:dataset:(urn:li:dataPlatform:sqlite,public.costs,PROD)",
Expand Down
Loading
Loading