acryl-datahub 0.15.0.1rc5__py3-none-any.whl → 0.15.0.1rc7__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -82,6 +82,7 @@ from datahub.ingestion.source_report.ingestion_stage import (
     LINEAGE_EXTRACTION,
     METADATA_EXTRACTION,
     QUERIES_EXTRACTION,
+    VIEW_PARSING,
 )
 from datahub.sql_parsing.sql_parsing_aggregator import SqlParsingAggregator
 from datahub.utilities.registries.domain_registry import DomainRegistry
@@ -103,7 +104,7 @@ logger: logging.Logger = logging.getLogger(__name__)
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 @capability(
     SourceCapability.LINEAGE_COARSE,
-    "Enabled by default, can be disabled via configuration `include_table_lineage` and `include_view_lineage`",
+    "Enabled by default, can be disabled via configuration `include_table_lineage`",
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
@@ -512,15 +513,14 @@ class SnowflakeV2Source(
         discovered_datasets = discovered_tables + discovered_views
 
         if self.config.use_queries_v2:
-            self.report.set_ingestion_stage("*", "View Parsing")
-            assert self.aggregator is not None
+            self.report.set_ingestion_stage("*", VIEW_PARSING)
             yield from auto_workunit(self.aggregator.gen_metadata())
 
             self.report.set_ingestion_stage("*", QUERIES_EXTRACTION)
 
             schema_resolver = self.aggregator._schema_resolver
 
-            queries_extractor: SnowflakeQueriesExtractor = SnowflakeQueriesExtractor(
+            queries_extractor = SnowflakeQueriesExtractor(
                 connection=self.connection,
                 config=SnowflakeQueriesExtractorConfig(
                     window=self.config,
@@ -546,13 +546,21 @@ class SnowflakeV2Source(
             queries_extractor.close()
 
         else:
-            if self.config.include_table_lineage and self.lineage_extractor:
+            if self.lineage_extractor:
                 self.report.set_ingestion_stage("*", LINEAGE_EXTRACTION)
-                yield from self.lineage_extractor.get_workunits(
+                self.lineage_extractor.add_time_based_lineage_to_aggregator(
                     discovered_tables=discovered_tables,
                     discovered_views=discovered_views,
                 )
 
+            # This would emit view and external table ddl lineage
+            # as well as query lineage via lineage_extractor
+            for mcp in self.aggregator.gen_metadata():
+                yield mcp.as_workunit()
+
+            if self.lineage_extractor:
+                self.lineage_extractor.update_state()
+
             if (
                 self.config.include_usage_stats or self.config.include_operational_stats
             ) and self.usage_extractor:

datahub/ingestion/source/tableau/tableau.py

@@ -35,7 +35,10 @@ from tableauserverclient import (
     SiteItem,
     TableauAuth,
 )
-from tableauserverclient.server.endpoint.exceptions import NonXMLResponseError
+from tableauserverclient.server.endpoint.exceptions import (
+    InternalServerError,
+    NonXMLResponseError,
+)
 from urllib3 import Retry
 
 import datahub.emitter.mce_builder as builder
@@ -618,6 +621,12 @@ class DatabaseTable:
                 self.parsed_columns = parsed_columns
 
 
+@dataclass
+class SiteIdContentUrl:
+    site_id: str
+    site_content_url: str
+
+
 class TableauSourceReport(StaleEntityRemovalSourceReport):
     get_all_datasources_query_failed: bool = False
     num_get_datasource_query_failures: int = 0
@@ -770,7 +779,6 @@ class TableauSource(StatefulIngestionSourceBase, TestableSource):
                         config=self.config,
                         ctx=self.ctx,
                         site=site,
-                        site_id=site.id,
                         report=self.report,
                         server=self.server,
                         platform=self.platform,
@@ -789,8 +797,11 @@ class TableauSource(StatefulIngestionSourceBase, TestableSource):
                 site_source = TableauSiteSource(
                     config=self.config,
                     ctx=self.ctx,
-                    site=site,
-                    site_id=self.server.site_id,
+                    site=site
+                    if site
+                    else SiteIdContentUrl(
+                        site_id=self.server.site_id, site_content_url=self.config.site
+                    ),
                     report=self.report,
                     server=self.server,
                     platform=self.platform,
@@ -823,8 +834,7 @@ class TableauSiteSource:
         self,
         config: TableauConfig,
         ctx: PipelineContext,
-        site: Optional[SiteItem],
-        site_id: Optional[str],
+        site: Union[SiteItem, SiteIdContentUrl],
         report: TableauSourceReport,
         server: Server,
         platform: str,
@@ -835,13 +845,18 @@ class TableauSiteSource:
         self.ctx: PipelineContext = ctx
         self.platform = platform
 
-        self.site: Optional[SiteItem] = site
-        if site_id is not None:
-            self.site_id: str = site_id
+        self.site: Optional[SiteItem] = None
+        if isinstance(site, SiteItem):
+            self.site = site
+            assert site.id is not None, "Site ID is required"
+            self.site_id = site.id
+            self.site_content_url = site.content_url
+        elif isinstance(site, SiteIdContentUrl):
+            self.site = None
+            self.site_id = site.site_id
+            self.site_content_url = site.site_content_url
         else:
-            assert self.site is not None, "site or site_id is required"
-            assert self.site.id is not None, "site_id is required when site is provided"
-            self.site_id = self.site.id
+            raise AssertionError("site or site id+content_url pair is required")
 
         self.database_tables: Dict[str, DatabaseTable] = {}
         self.tableau_stat_registry: Dict[str, UsageStat] = {}
@@ -895,16 +910,14 @@ class TableauSiteSource:
         # datasets also have the env in the browse path
         return f"/{self.config.env.lower()}{self.no_env_browse_prefix}"
 
-    def _re_authenticate(self):
+    def _re_authenticate(self) -> None:
+        self.report.info(
+            message="Re-authenticating to Tableau",
+            context=f"site='{self.site_content_url}'",
+        )
         # Sign-in again may not be enough because Tableau sometimes caches invalid sessions
         # so we need to recreate the Tableau Server object
-        self.server = self.config.make_tableau_client(self.site_id)
-
-    @property
-    def site_content_url(self) -> Optional[str]:
-        if self.site and self.site.content_url:
-            return self.site.content_url
-        return None
+        self.server = self.config.make_tableau_client(self.site_content_url)
 
     def _populate_usage_stat_registry(self) -> None:
         if self.server is None:
@@ -1196,6 +1209,24 @@ class TableauSiteSource:
                 retry_on_auth_error=False,
                 retries_remaining=retries_remaining - 1,
             )
+
+        except InternalServerError as ise:
+            # In some cases Tableau Server returns 504 error, which is a timeout error, so it worths to retry.
+            if ise.code == 504:
+                if retries_remaining <= 0:
+                    raise ise
+                return self.get_connection_object_page(
+                    query=query,
+                    connection_type=connection_type,
+                    query_filter=query_filter,
+                    fetch_size=fetch_size,
+                    current_cursor=current_cursor,
+                    retry_on_auth_error=False,
+                    retries_remaining=retries_remaining - 1,
+                )
+            else:
+                raise ise
+
         except OSError:
             # In tableauseverclient 0.26 (which was yanked and released in 0.28 on 2023-10-04),
             # the request logic was changed to use threads.

datahub/ingestion/source_report/ingestion_stage.py

@@ -15,6 +15,7 @@ USAGE_EXTRACTION_INGESTION = "Usage Extraction Ingestion"
 USAGE_EXTRACTION_OPERATIONAL_STATS = "Usage Extraction Operational Stats"
 USAGE_EXTRACTION_USAGE_AGGREGATION = "Usage Extraction Usage Aggregation"
 EXTERNAL_TABLE_DDL_LINEAGE = "External table DDL Lineage"
+VIEW_PARSING = "View Parsing"
 QUERIES_EXTRACTION = "Queries Extraction"
 PROFILING = "Profiling"
 

datahub/metadata/_schema_classes.py

@@ -14737,7 +14737,7 @@ class DataProcessInstanceKeyClass(_Aspect):
 
 
     ASPECT_NAME = 'dataProcessInstanceKey'
-    ASPECT_INFO = {'keyForEntity': 'dataProcessInstance', 'entityCategory': '_unset_', 'entityAspects': ['dataProcessInstanceInput', 'dataProcessInstanceOutput', 'dataProcessInstanceProperties', 'dataProcessInstanceRelationships', 'dataProcessInstanceRunEvent', 'status', 'testResults'], 'entityDoc': 'DataProcessInstance represents an instance of a datajob/jobflow run'}
+    ASPECT_INFO = {'keyForEntity': 'dataProcessInstance', 'entityCategory': '_unset_', 'entityAspects': ['dataProcessInstanceInput', 'dataProcessInstanceOutput', 'dataProcessInstanceProperties', 'dataProcessInstanceRelationships', 'dataProcessInstanceRunEvent', 'status', 'testResults', 'dataPlatformInstance', 'subTypes', 'container', 'mlTrainingRunProperties'], 'entityDoc': 'DataProcessInstance represents an instance of a datajob/jobflow run'}
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.metadata.key.DataProcessInstanceKey")
 
     def __init__(self,
@@ -17715,8 +17715,12 @@ class MLModelGroupPropertiesClass(_Aspect):
 
     def __init__(self,
         customProperties: Optional[Dict[str, str]]=None,
+        name: Union[None, str]=None,
         description: Union[None, str]=None,
         createdAt: Union[None, int]=None,
+        created: Union[None, "TimeStampClass"]=None,
+        lastModified: Union[None, "TimeStampClass"]=None,
+        trainingJobs: Union[None, List[str]]=None,
         version: Union[None, "VersionTagClass"]=None,
     ):
         super().__init__()
@@ -17726,14 +17730,22 @@ class MLModelGroupPropertiesClass(_Aspect):
             self.customProperties = dict()
         else:
             self.customProperties = customProperties
+        self.name = name
         self.description = description
         self.createdAt = createdAt
+        self.created = created
+        self.lastModified = lastModified
+        self.trainingJobs = trainingJobs
         self.version = version
     
     def _restore_defaults(self) -> None:
         self.customProperties = dict()
+        self.name = self.RECORD_SCHEMA.fields_dict["name"].default
         self.description = self.RECORD_SCHEMA.fields_dict["description"].default
         self.createdAt = self.RECORD_SCHEMA.fields_dict["createdAt"].default
+        self.created = self.RECORD_SCHEMA.fields_dict["created"].default
+        self.lastModified = self.RECORD_SCHEMA.fields_dict["lastModified"].default
+        self.trainingJobs = self.RECORD_SCHEMA.fields_dict["trainingJobs"].default
         self.version = self.RECORD_SCHEMA.fields_dict["version"].default
     
     
@@ -17747,6 +17759,16 @@ class MLModelGroupPropertiesClass(_Aspect):
         self._inner_dict['customProperties'] = value
     
     
+    @property
+    def name(self) -> Union[None, str]:
+        """Display name of the MLModelGroup"""
+        return self._inner_dict.get('name')  # type: ignore
+    
+    @name.setter
+    def name(self, value: Union[None, str]) -> None:
+        self._inner_dict['name'] = value
+    
+    
     @property
     def description(self) -> Union[None, str]:
         """Documentation of the MLModelGroup"""
@@ -17767,6 +17789,36 @@ class MLModelGroupPropertiesClass(_Aspect):
         self._inner_dict['createdAt'] = value
     
     
+    @property
+    def created(self) -> Union[None, "TimeStampClass"]:
+        """Time and Actor who created the MLModelGroup"""
+        return self._inner_dict.get('created')  # type: ignore
+    
+    @created.setter
+    def created(self, value: Union[None, "TimeStampClass"]) -> None:
+        self._inner_dict['created'] = value
+    
+    
+    @property
+    def lastModified(self) -> Union[None, "TimeStampClass"]:
+        """Date when the MLModelGroup was last modified"""
+        return self._inner_dict.get('lastModified')  # type: ignore
+    
+    @lastModified.setter
+    def lastModified(self, value: Union[None, "TimeStampClass"]) -> None:
+        self._inner_dict['lastModified'] = value
+    
+    
+    @property
+    def trainingJobs(self) -> Union[None, List[str]]:
+        """List of jobs (if any) used to train the model group. Visible in Lineage."""
+        return self._inner_dict.get('trainingJobs')  # type: ignore
+    
+    @trainingJobs.setter
+    def trainingJobs(self, value: Union[None, List[str]]) -> None:
+        self._inner_dict['trainingJobs'] = value
+    
+    
     @property
     def version(self) -> Union[None, "VersionTagClass"]:
         """Version of the MLModelGroup"""
@@ -17788,8 +17840,11 @@ class MLModelPropertiesClass(_Aspect):
     def __init__(self,
         customProperties: Optional[Dict[str, str]]=None,
         externalUrl: Union[None, str]=None,
+        name: Union[None, str]=None,
         description: Union[None, str]=None,
         date: Union[None, int]=None,
+        created: Union[None, "TimeStampClass"]=None,
+        lastModified: Union[None, "TimeStampClass"]=None,
         version: Union[None, "VersionTagClass"]=None,
         type: Union[None, str]=None,
         hyperParameters: Union[None, Dict[str, Union[str, int, float, float, bool]]]=None,
@@ -17811,8 +17866,11 @@ class MLModelPropertiesClass(_Aspect):
         else:
             self.customProperties = customProperties
         self.externalUrl = externalUrl
+        self.name = name
         self.description = description
         self.date = date
+        self.created = created
+        self.lastModified = lastModified
         self.version = version
         self.type = type
         self.hyperParameters = hyperParameters
@@ -17833,8 +17891,11 @@ class MLModelPropertiesClass(_Aspect):
     def _restore_defaults(self) -> None:
         self.customProperties = dict()
         self.externalUrl = self.RECORD_SCHEMA.fields_dict["externalUrl"].default
+        self.name = self.RECORD_SCHEMA.fields_dict["name"].default
         self.description = self.RECORD_SCHEMA.fields_dict["description"].default
         self.date = self.RECORD_SCHEMA.fields_dict["date"].default
+        self.created = self.RECORD_SCHEMA.fields_dict["created"].default
+        self.lastModified = self.RECORD_SCHEMA.fields_dict["lastModified"].default
         self.version = self.RECORD_SCHEMA.fields_dict["version"].default
         self.type = self.RECORD_SCHEMA.fields_dict["type"].default
         self.hyperParameters = self.RECORD_SCHEMA.fields_dict["hyperParameters"].default
@@ -17869,6 +17930,16 @@ class MLModelPropertiesClass(_Aspect):
         self._inner_dict['externalUrl'] = value
     
     
+    @property
+    def name(self) -> Union[None, str]:
+        """Display name of the MLModel"""
+        return self._inner_dict.get('name')  # type: ignore
+    
+    @name.setter
+    def name(self, value: Union[None, str]) -> None:
+        self._inner_dict['name'] = value
+    
+    
     @property
     def description(self) -> Union[None, str]:
         """Documentation of the MLModel"""
@@ -17889,6 +17960,26 @@ class MLModelPropertiesClass(_Aspect):
         self._inner_dict['date'] = value
     
     
+    @property
+    def created(self) -> Union[None, "TimeStampClass"]:
+        """Audit stamp containing who created this and when"""
+        return self._inner_dict.get('created')  # type: ignore
+    
+    @created.setter
+    def created(self, value: Union[None, "TimeStampClass"]) -> None:
+        self._inner_dict['created'] = value
+    
+    
+    @property
+    def lastModified(self) -> Union[None, "TimeStampClass"]:
+        """Date when the MLModel was last modified"""
+        return self._inner_dict.get('lastModified')  # type: ignore
+    
+    @lastModified.setter
+    def lastModified(self, value: Union[None, "TimeStampClass"]) -> None:
+        self._inner_dict['lastModified'] = value
+    
+    
     @property
     def version(self) -> Union[None, "VersionTagClass"]:
         """Version of the MLModel"""
@@ -17983,7 +18074,7 @@ class MLModelPropertiesClass(_Aspect):
     
     @property
     def trainingJobs(self) -> Union[None, List[str]]:
-        """List of jobs (if any) used to train the model"""
+        """List of jobs (if any) used to train the model. Visible in Lineage. Note that ML Models can also be specified as the output of a specific Data Process Instances (runs) via the DataProcessInstanceOutputs aspect."""
         return self._inner_dict.get('trainingJobs')  # type: ignore
     
     @trainingJobs.setter
@@ -18096,6 +18187,104 @@ class MLPrimaryKeyPropertiesClass(_Aspect):
         self._inner_dict['sources'] = value
     
     
+class MLTrainingRunPropertiesClass(_Aspect):
+    """The inputs and outputs of this training run"""
+
+
+    ASPECT_NAME = 'mlTrainingRunProperties'
+    ASPECT_INFO = {}
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.ml.metadata.MLTrainingRunProperties")
+
+    def __init__(self,
+        customProperties: Optional[Dict[str, str]]=None,
+        externalUrl: Union[None, str]=None,
+        id: Union[None, str]=None,
+        outputUrls: Union[None, List[str]]=None,
+        hyperParams: Union[None, List["MLHyperParamClass"]]=None,
+        trainingMetrics: Union[None, List["MLMetricClass"]]=None,
+    ):
+        super().__init__()
+        
+        if customProperties is None:
+            # default: {}
+            self.customProperties = dict()
+        else:
+            self.customProperties = customProperties
+        self.externalUrl = externalUrl
+        self.id = id
+        self.outputUrls = outputUrls
+        self.hyperParams = hyperParams
+        self.trainingMetrics = trainingMetrics
+    
+    def _restore_defaults(self) -> None:
+        self.customProperties = dict()
+        self.externalUrl = self.RECORD_SCHEMA.fields_dict["externalUrl"].default
+        self.id = self.RECORD_SCHEMA.fields_dict["id"].default
+        self.outputUrls = self.RECORD_SCHEMA.fields_dict["outputUrls"].default
+        self.hyperParams = self.RECORD_SCHEMA.fields_dict["hyperParams"].default
+        self.trainingMetrics = self.RECORD_SCHEMA.fields_dict["trainingMetrics"].default
+    
+    
+    @property
+    def customProperties(self) -> Dict[str, str]:
+        """Custom property bag."""
+        return self._inner_dict.get('customProperties')  # type: ignore
+    
+    @customProperties.setter
+    def customProperties(self, value: Dict[str, str]) -> None:
+        self._inner_dict['customProperties'] = value
+    
+    
+    @property
+    def externalUrl(self) -> Union[None, str]:
+        """URL where the reference exist"""
+        return self._inner_dict.get('externalUrl')  # type: ignore
+    
+    @externalUrl.setter
+    def externalUrl(self, value: Union[None, str]) -> None:
+        self._inner_dict['externalUrl'] = value
+    
+    
+    @property
+    def id(self) -> Union[None, str]:
+        """Run Id of the ML Training Run"""
+        return self._inner_dict.get('id')  # type: ignore
+    
+    @id.setter
+    def id(self, value: Union[None, str]) -> None:
+        self._inner_dict['id'] = value
+    
+    
+    @property
+    def outputUrls(self) -> Union[None, List[str]]:
+        """List of URLs for the Outputs of the ML Training Run"""
+        return self._inner_dict.get('outputUrls')  # type: ignore
+    
+    @outputUrls.setter
+    def outputUrls(self, value: Union[None, List[str]]) -> None:
+        self._inner_dict['outputUrls'] = value
+    
+    
+    @property
+    def hyperParams(self) -> Union[None, List["MLHyperParamClass"]]:
+        """Hyperparameters of the ML Training Run"""
+        return self._inner_dict.get('hyperParams')  # type: ignore
+    
+    @hyperParams.setter
+    def hyperParams(self, value: Union[None, List["MLHyperParamClass"]]) -> None:
+        self._inner_dict['hyperParams'] = value
+    
+    
+    @property
+    def trainingMetrics(self) -> Union[None, List["MLMetricClass"]]:
+        """Metrics of the ML Training Run"""
+        return self._inner_dict.get('trainingMetrics')  # type: ignore
+    
+    @trainingMetrics.setter
+    def trainingMetrics(self, value: Union[None, List["MLMetricClass"]]) -> None:
+        self._inner_dict['trainingMetrics'] = value
+    
+    
 class MetricsClass(_Aspect):
     """Metrics to be featured for the MLModel."""
 
@@ -24791,6 +24980,7 @@ __SCHEMA_TYPES = {
     'com.linkedin.pegasus2avro.ml.metadata.MLModelGroupProperties': MLModelGroupPropertiesClass,
     'com.linkedin.pegasus2avro.ml.metadata.MLModelProperties': MLModelPropertiesClass,
     'com.linkedin.pegasus2avro.ml.metadata.MLPrimaryKeyProperties': MLPrimaryKeyPropertiesClass,
+    'com.linkedin.pegasus2avro.ml.metadata.MLTrainingRunProperties': MLTrainingRunPropertiesClass,
     'com.linkedin.pegasus2avro.ml.metadata.Metrics': MetricsClass,
     'com.linkedin.pegasus2avro.ml.metadata.QuantitativeAnalyses': QuantitativeAnalysesClass,
     'com.linkedin.pegasus2avro.ml.metadata.SourceCode': SourceCodeClass,
@@ -25257,6 +25447,7 @@ __SCHEMA_TYPES = {
     'MLModelGroupProperties': MLModelGroupPropertiesClass,
     'MLModelProperties': MLModelPropertiesClass,
     'MLPrimaryKeyProperties': MLPrimaryKeyPropertiesClass,
+    'MLTrainingRunProperties': MLTrainingRunPropertiesClass,
     'Metrics': MetricsClass,
     'QuantitativeAnalyses': QuantitativeAnalysesClass,
     'SourceCode': SourceCodeClass,
@@ -25473,6 +25664,7 @@ ASPECT_CLASSES: List[Type[_Aspect]] = [
     MetricsClass,
     EditableMLPrimaryKeyPropertiesClass,
     SourceCodeClass,
+    MLTrainingRunPropertiesClass,
     EthicalConsiderationsClass,
     MLPrimaryKeyPropertiesClass,
     MLModelFactorPromptsClass,
@@ -25686,6 +25878,7 @@ class AspectBag(TypedDict, total=False):
     mlModelMetrics: MetricsClass
     editableMlPrimaryKeyProperties: EditableMLPrimaryKeyPropertiesClass
     sourceCode: SourceCodeClass
+    mlTrainingRunProperties: MLTrainingRunPropertiesClass
     mlModelEthicalConsiderations: EthicalConsiderationsClass
     mlPrimaryKeyProperties: MLPrimaryKeyPropertiesClass
     mlModelFactorPrompts: MLModelFactorPromptsClass

datahub/metadata/com/linkedin/pegasus2avro/ml/metadata/__init__.py

@@ -30,6 +30,7 @@ from ......schema_classes import MLModelFactorsClass
 from ......schema_classes import MLModelGroupPropertiesClass
 from ......schema_classes import MLModelPropertiesClass
 from ......schema_classes import MLPrimaryKeyPropertiesClass
+from ......schema_classes import MLTrainingRunPropertiesClass
 from ......schema_classes import MetricsClass
 from ......schema_classes import QuantitativeAnalysesClass
 from ......schema_classes import SourceCodeClass
@@ -61,6 +62,7 @@ MLModelFactors = MLModelFactorsClass
 MLModelGroupProperties = MLModelGroupPropertiesClass
 MLModelProperties = MLModelPropertiesClass
 MLPrimaryKeyProperties = MLPrimaryKeyPropertiesClass
+MLTrainingRunProperties = MLTrainingRunPropertiesClass
 Metrics = MetricsClass
 QuantitativeAnalyses = QuantitativeAnalysesClass
 SourceCode = SourceCodeClass