feat(ingest/unity): Add usage extraction; add TableReference #7910
Conversation
github-actions
bot
added
the
ingestion
| def spark_sql_parser(self): | ||
| """Lazily initializes the Spark SQL parser.""" | ||
| if self._spark_sql_parser is None: | ||
| spark_context = pyspark.SparkContext.getOrCreate() |
There was a problem hiding this comment.
Do we install pyspark dependency for unity-catalog ? Its not explicitly mentioned in setup.py. Is it indirectly installed for databricks-cli ?
There was a problem hiding this comment.
Ah good catch, it's part of delta lake but not databricks
|
|
||
|
|
||
| class UnityCatalogSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigMixin): | ||
| class UnityCatalogSourceConfig( | ||
| StatefulIngestionConfigBase, BaseUsageConfig, DatasetSourceConfigMixin |
There was a problem hiding this comment.
BaseUsageConfig has a lot of fields - if not all of them are supported, then we should change the base class that we're using
There was a problem hiding this comment.
We use most of them, as they get used by usage_common. The only one that doesn't get used is include_read_operational_stats which seems to only be used by bigquery... so if anything I say we remove it from this common config and put it only in the bigquery one.
| ) | ||
| table.upstreams.setdefault(table_ref, {}).setdefault( | ||
| column.name, [] | ||
| ).append(item["name"]) |
|
|
||
|
|
||
| TableMap = Dict[str, List[TableReference]] | ||
| T = TypeVar("T", bound=object) |
There was a problem hiding this comment.
don't think the bound is doing anything here
| ) | ||
|
|
||
| if not self.config.table_pattern.allowed(filter_table_name): | ||
| if not self.config.table_pattern.allowed(table.ref.qualified_table_name): |
There was a problem hiding this comment.
the filter_table_name isn't quite the fully qualified table name, since it doesn't include the metastore name
There was a problem hiding this comment.
Yeah, so I made TableReference.qualified_table_name not include metastore name, while TableReference.__str__ includes metastore name. I'm realizing that naming can be confusing though. Any ideas on a better one?
|
|
||
| def _parse_query_via_lineage_runner(self, query: str) -> Optional[StringTableInfo]: | ||
| try: | ||
| runner = LineageRunner(query) |
There was a problem hiding this comment.
not a huge fan of the fact that this LineageRunner stuff is copy-pasted across our codebase
Ideally we'd go through a unified parser interface, where you can select the underlying parser(s) to use with an enum or something
There was a problem hiding this comment.
don't need to fix it here though
There was a problem hiding this comment.
We have a our SqlLineageSQLParser which creates SqlLineageSQLParserImpl which eventually calls LineageAnalyzer().analyze, but there was a lot of logic in there which didn't seem necessary. I thought overall this was simpler. Agree that we should standardize this as much as possible, but I think it'll take a bit of work
| spark_context = pyspark.SparkContext.getOrCreate() | ||
| spark_session = pyspark.sql.SparkSession(spark_context) | ||
| self._spark_sql_parser = ( | ||
| spark_session._jsparkSession.sessionState().sqlParser() |
Adds unity catalog usage extraction
Problems encountered:
Uses sqllineage parser and a spark sql parser as fallback.
Refactors:
UsageAggregatorclass to usage_common, as I've seen this same logic multiple times.user_urn_builderin usage_common as not all unity users are emails. We create emails with a defaultemail_domainconfig in other connectors like redshift and snowflake, which seems unnecessary now?TableReferencefor unity catalog and adds it to theTabledataclass, for managing string references to tables. Replaces logic, especially in lineage extraction, with these referencesgen_dataset_urnandgen_user_urnon unity source to reduce duplicate codeproxy.pyinto implementation and typesChecklist