feat(ingest/bigquery): support column-level lineage #8382
Conversation
There was a problem hiding this comment.
Overall looks good to me! If we're looking to ship this soon for a client, I don't remember any changes that /need/ to happen, so I think we can push and address followups after. I got stuck on the recursive temp-table-removal process for a bit, so I have some naming suggestions there but not sure the best way to make that clearer
metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py
Outdated
Show resolved
Hide resolved
metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/lineage.py
Outdated
Show resolved
Hide resolved
| auditStamp=datetime.now(timezone.utc), | ||
| type=DatasetLineageTypeClass.VIEW, | ||
| for view, view_definition in self.view_definitions[project_id].items(): | ||
| raw_view_lineage = sqlglot_lineage( |
There was a problem hiding this comment.
Do we think this sql parser is strictly better than the old one? A fallback to the old one might be safer, but honestly I would like to err on the side of code velocity over safety (in cases like these) so I am fine with this as is
EDIT: That being said, a function like below might be nice, in case we support multiple parsers or adjust the call signature of sqlglot_lineage:
def parse_lineage(self, query: str, project_id: str):
return sqlglot_lineage(query, self.platform, self.sql_parser_schema_resolver, project_id)
There was a problem hiding this comment.
I initially wanted to do that
However the sql_parser_schema_resolver is only available in the main BigquerySource, and not in the BigqueryLineage class.
Ideally we create a BigqueryContext object with the schema resolver, urn generation utils, etc and pass that around everywhere
| lineage[view] = set( | ||
| make_lineage_edges_from_parsing_result( | ||
| raw_view_lineage, | ||
| audit_stamp=ts, | ||
| lineage_type=DatasetLineageTypeClass.VIEW, | ||
| ) | ||
| for table in upstream_tables | ||
| } | ||
| ) |
There was a problem hiding this comment.
I see we always cast the result of make_lineage_edges_from_parsing_result to set, could we just do that in make_lineage_edges_from_parsing_result instead?
metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/lineage.py
Outdated
Show resolved
Hide resolved
metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/lineage.py
Outdated
Show resolved
Hide resolved
| @@ -629,23 +785,26 @@ def calculate_lineage_for_project( | |||
| def get_lineage_for_table( | |||
| self, | |||
| bq_table: BigQueryTableRef, | |||
| bq_table_urn: str, | |||
There was a problem hiding this comment.
Doesn't have to be done in this PR, but I'd like to pass into the lineage extractor dataset_urn_builder: Callable[[BigQueryTableRef], str] like I do with the bq usage extractor, so that we always use the same method to generate dataset urns. Then, I think this won't be necessary (and we can remove the mce_builder call below in favor of it too)
| self.sql_parser_schema_resolver = SchemaResolver( | ||
| platform=self.platform, env=self.config.env | ||
| ) |
There was a problem hiding this comment.
Ideally, I think we'd implement SchemaResolver.close and call it when this source is closed and when the datahub graph is closed, rather than relying on garbage collection, which I think is the case right now
|
The CI issues will be fixed by #8393, so merging through. |
Changes stacked on #8334.
Checklist