acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/powerbi/m_query/parser.py

@@ -1,13 +1,13 @@
 import functools
 import importlib.resources as pkg_resource
 import logging
-import os
 from typing import Dict, List, Optional
 
 import lark
 from lark import Lark, Tree
 
 import datahub.ingestion.source.powerbi.m_query.data_classes
+from datahub.configuration.env_vars import get_powerbi_m_query_parse_timeout
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.source.powerbi.config import (
     PowerBiDashboardSourceConfig,
@@ -25,7 +25,7 @@ from datahub.utilities.threading_timeout import TimeoutException, threading_time
 
 logger = logging.getLogger(__name__)
 
-_M_QUERY_PARSE_TIMEOUT = int(os.getenv("DATAHUB_POWERBI_M_QUERY_PARSE_TIMEOUT", 60))
+_M_QUERY_PARSE_TIMEOUT = get_powerbi_m_query_parse_timeout()
 
 
 @functools.lru_cache(maxsize=1)

datahub/ingestion/source/powerbi/m_query/pattern_handler.py

@@ -3,7 +3,9 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Dict, List, Optional, Tuple, Type, cast
 
+import sqlglot
 from lark import Tree
+from sqlglot import ParseError, expressions as exp
 
 from datahub.configuration.source_common import PlatformDetail
 from datahub.emitter import mce_builder as builder
@@ -209,15 +211,34 @@ class AbstractLineage(ABC):
 
         return None
 
+    @staticmethod
+    def is_sql_query(query: Optional[str]) -> bool:
+        if not query:
+            return False
+        query = native_sql_parser.remove_special_characters(query)
+        try:
+            expression = sqlglot.parse_one(query)
+            return isinstance(expression, exp.Select)
+        except (ParseError, Exception):
+            logger.debug(f"Failed to parse query as SQL: {query}")
+            return False
+
     def parse_custom_sql(
-        self, query: str, server: str, database: Optional[str], schema: Optional[str]
+        self,
+        query: str,
+        server: str,
+        database: Optional[str],
+        schema: Optional[str],
+        platform_pair: Optional[DataPlatformPair] = None,
     ) -> Lineage:
         dataplatform_tables: List[DataPlatformTable] = []
+        if not platform_pair:
+            platform_pair = self.get_platform_pair()
 
         platform_detail: PlatformDetail = (
             self.platform_instance_resolver.get_platform_instance(
                 PowerBIPlatformDetail(
-                    data_platform_pair=self.get_platform_pair(),
+                    data_platform_pair=platform_pair,
                     data_platform_server=server,
                 )
             )
@@ -231,7 +252,7 @@ class AbstractLineage(ABC):
             native_sql_parser.parse_custom_sql(
                 ctx=self.ctx,
                 query=query,
-                platform=self.get_platform_pair().datahub_data_platform_name,
+                platform=platform_pair.datahub_data_platform_name,
                 platform_instance=platform_detail.platform_instance,
                 env=platform_detail.env,
                 database=database,
@@ -258,7 +279,7 @@ class AbstractLineage(ABC):
         for urn in parsed_result.in_tables:
             dataplatform_tables.append(
                 DataPlatformTable(
-                    data_platform_pair=self.get_platform_pair(),
+                    data_platform_pair=platform_pair,
                     urn=urn,
                 )
             )
@@ -956,7 +977,7 @@ class OdbcLineage(AbstractLineage):
             f"data-access function detail {data_access_func_detail}"
         )
 
-        connect_string, _ = self.get_db_detail_from_argument(
+        connect_string, query = self.get_db_detail_from_argument(
             data_access_func_detail.arg_list
         )
 
@@ -972,12 +993,19 @@ class OdbcLineage(AbstractLineage):
         data_platform, powerbi_platform = extract_platform(connect_string)
         server_name = extract_server(connect_string)
 
+        dsn = extract_dsn(connect_string)
+        if not dsn:
+            self.reporter.warning(
+                title="Can not determine ODBC DSN",
+                message="Can not extract DSN from ODBC connect string. Skipping Lineage creation.",
+                context=f"table-name={self.table.full_name}, connect-string={connect_string}",
+            )
+            return Lineage.empty()
+        logger.debug(f"Extracted DSN: {dsn}")
+
         if not data_platform:
-            dsn = extract_dsn(connect_string)
-            if dsn:
-                logger.debug(f"Extracted DSN: {dsn}")
-                server_name = dsn
-            if dsn and self.config.dsn_to_platform_name:
+            server_name = dsn
+            if self.config.dsn_to_platform_name:
                 logger.debug(f"Attempting to map DSN {dsn} to platform")
                 name = self.config.dsn_to_platform_name.get(dsn)
                 if name:
@@ -1006,6 +1034,63 @@ class OdbcLineage(AbstractLineage):
         elif not server_name:
             server_name = "unknown"
 
+        if self.is_sql_query(query):
+            return self.query_lineage(query, platform_pair, server_name, dsn)
+        else:
+            return self.expression_lineage(
+                data_access_func_detail, data_platform, platform_pair, server_name
+            )
+
+    def query_lineage(
+        self,
+        query: Optional[str],
+        platform_pair: DataPlatformPair,
+        server_name: str,
+        dsn: str,
+    ) -> Lineage:
+        database = None
+        schema = None
+
+        if not query:
+            # query should never be None as it is checked before calling this function.
+            # however, we need to check just in case.
+            self.reporter.warning(
+                title="ODBC Query is null",
+                message="No SQL to parse. Skipping Lineage creation.",
+                context=f"table-name={self.table.full_name}",
+            )
+            return Lineage.empty()
+
+        if self.config.dsn_to_database_schema:
+            value = self.config.dsn_to_database_schema.get(dsn)
+            if value:
+                parts = value.split(".")
+                if len(parts) == 1:
+                    database = parts[0]
+                elif len(parts) == 2:
+                    database = parts[0]
+                    schema = parts[1]
+
+        logger.debug(
+            f"ODBC query processing: dsn={dsn} mapped to database={database}, schema={schema}"
+        )
+        result = self.parse_custom_sql(
+            query=query,
+            server=server_name,
+            database=database,
+            schema=schema,
+            platform_pair=platform_pair,
+        )
+        logger.debug(f"ODBC query lineage generated {len(result.upstreams)} upstreams")
+        return result
+
+    def expression_lineage(
+        self,
+        data_access_func_detail: DataAccessFunctionDetail,
+        data_platform: str,
+        platform_pair: DataPlatformPair,
+        server_name: str,
+    ) -> Lineage:
         database_name = None
         schema_name = None
         table_name = None
@@ -1144,6 +1229,11 @@ class SupportedPattern(Enum):
         FunctionName.ODBC_DATA_ACCESS,
     )
 
+    ODBC_QUERY = (
+        OdbcLineage,
+        FunctionName.ODBC_QUERY,
+    )
+
     def handler(self) -> Type[AbstractLineage]:
         return self.value[0]
 

datahub/ingestion/source/powerbi/powerbi.py

@@ -40,6 +40,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import (
     BIAssetSubTypes,
     BIContainerSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.powerbi.config import (
     Constant,
@@ -294,8 +295,6 @@ class Mapper:
             logger.debug(f"Dataset urn = {ds_urn} and its lineage = {upstream_lineage}")
 
             mcp = MetadataChangeProposalWrapper(
-                entityType=Constant.DATASET,
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=ds_urn,
                 aspect=upstream_lineage_class,
             )
@@ -538,9 +537,7 @@ class Mapper:
         profile.columnCount = table.column_count
 
         mcp = MetadataChangeProposalWrapper(
-            entityType="dataset",
             entityUrn=ds_urn,
-            aspectName="datasetProfile",
             aspect=profile,
         )
         dataset_mcps.append(mcp)
@@ -796,7 +793,6 @@ class Mapper:
             guid=container_key.guid(),
         )
         mcp = MetadataChangeProposalWrapper(
-            changeType=ChangeTypeClass.UPSERT,
             entityUrn=entity_urn,
             aspect=ContainerClass(container=f"{container_urn}"),
         )
@@ -1231,7 +1227,14 @@ class Mapper:
 @platform_name("PowerBI")
 @config_class(PowerBiDashboardSourceConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.POWERBI_WORKSPACE,
+        SourceCapabilityModifier.POWERBI_DATASET,
+    ],
+)
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 @capability(SourceCapability.OWNERSHIP, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
@@ -1253,6 +1256,7 @@ class Mapper:
     SourceCapability.DATA_PROFILING,
     "Optionally enabled via configuration profiling.enabled",
 )
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class PowerBiDashboardSource(StatefulIngestionSourceBase, TestableSource):
     """
     This plugin extracts the following:

datahub/ingestion/source/powerbi/rest_api_wrapper/powerbi_api.py

@@ -673,7 +673,6 @@ class PowerBiAPI:
         fill_dashboard_tags()
         self._fill_independent_datasets(workspace=workspace)
 
-    # flake8: noqa: C901
     def fill_workspaces(
         self, workspaces: List[Workspace], reporter: PowerBiDashboardSourceReport
     ) -> Iterable[Workspace]:

datahub/ingestion/source/powerbi_report_server/report_server.py

@@ -52,7 +52,6 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
 from datahub.metadata.com.linkedin.pegasus2avro.common import ChangeAuditStamps
 from datahub.metadata.schema_classes import (
     BrowsePathsClass,
-    ChangeTypeClass,
     CorpUserInfoClass,
     CorpUserKeyClass,
     DashboardInfoClass,
@@ -243,20 +242,14 @@ class Mapper:
 
     @staticmethod
     def new_mcp(
-        entity_type,
         entity_urn,
-        aspect_name,
         aspect,
-        change_type=ChangeTypeClass.UPSERT,
     ):
         """
         Create MCP
         """
         return MetadataChangeProposalWrapper(
-            entityType=entity_type,
-            changeType=change_type,
             entityUrn=entity_urn,
-            aspectName=aspect_name,
             aspect=aspect,
         )
 
@@ -343,17 +336,13 @@ class Mapper:
         )
 
         info_mcp = self.new_mcp(
-            entity_type=Constant.DASHBOARD,
             entity_urn=dashboard_urn,
-            aspect_name=Constant.DASHBOARD_INFO,
             aspect=dashboard_info_cls,
         )
 
         # removed status mcp
         removed_status_mcp = self.new_mcp(
-            entity_type=Constant.DASHBOARD,
             entity_urn=dashboard_urn,
-            aspect_name=Constant.STATUS,
             aspect=StatusClass(removed=False),
         )
 
@@ -365,9 +354,7 @@ class Mapper:
 
         # Dashboard key
         dashboard_key_mcp = self.new_mcp(
-            entity_type=Constant.DASHBOARD,
             entity_urn=dashboard_urn,
-            aspect_name=Constant.DASHBOARD_KEY,
             aspect=dashboard_key_cls,
         )
 
@@ -378,9 +365,7 @@ class Mapper:
         ownership = OwnershipClass(owners=owners)
         # Dashboard owner MCP
         owner_mcp = self.new_mcp(
-            entity_type=Constant.DASHBOARD,
             entity_urn=dashboard_urn,
-            aspect_name=Constant.OWNERSHIP,
             aspect=ownership,
         )
 
@@ -396,9 +381,7 @@ class Mapper:
             ]
         )
         browse_path_mcp = self.new_mcp(
-            entity_type=Constant.DASHBOARD,
             entity_urn=dashboard_urn,
-            aspect_name=Constant.BROWSERPATH,
             aspect=browse_path,
         )
 
@@ -429,27 +412,21 @@ class Mapper:
             )
 
             info_mcp = self.new_mcp(
-                entity_type=Constant.CORP_USER,
                 entity_urn=user_urn,
-                aspect_name=Constant.CORP_USER_INFO,
                 aspect=user_info_instance,
             )
             user_mcps.append(info_mcp)
 
             # removed status mcp
             status_mcp = self.new_mcp(
-                entity_type=Constant.CORP_USER,
                 entity_urn=user_urn,
-                aspect_name=Constant.STATUS,
                 aspect=StatusClass(removed=False),
             )
             user_mcps.append(status_mcp)
             user_key = CorpUserKeyClass(username=user.username)
 
             user_key_mcp = self.new_mcp(
-                entity_type=Constant.CORP_USER,
                 entity_urn=user_urn,
-                aspect_name=Constant.CORP_USER_KEY,
                 aspect=user_key,
             )
             user_mcps.append(user_key_mcp)

datahub/ingestion/source/powerbi_report_server/report_server_domain.py

@@ -27,10 +27,8 @@ class CatalogItem(BaseModel):
     is_favorite: bool = Field(alias="IsFavorite")
     user_info: Any = Field(None, alias="UserInfo")
     display_name: Optional[str] = Field(None, alias="DisplayName")
-    has_data_sources: bool = Field(default=False, alias="HasDataSources")
-    data_sources: Optional[List["DataSource"]] = Field(
-        default_factory=list, alias="DataSources"
-    )
+    has_data_sources: bool = Field(False, alias="HasDataSources")
+    data_sources: Optional[List["DataSource"]] = Field(None, alias="DataSources")
 
     @validator("display_name", always=True)
     def validate_diplay_name(cls, value, values):

datahub/ingestion/source/preset.py

@@ -2,7 +2,7 @@ import logging
 from typing import Dict, Optional
 
 import requests
-from pydantic.class_validators import root_validator, validator
+from pydantic import root_validator, validator
 from pydantic.fields import Field
 
 from datahub.emitter.mce_builder import DEFAULT_ENV
@@ -69,9 +69,9 @@ class PresetConfig(SupersetConfig):
 
 @platform_name("Preset")
 @config_class(PresetConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.CERTIFIED)
 @capability(
-    SourceCapability.DELETION_DETECTION, "Optionally enabled via stateful_ingestion"
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 class PresetSource(SupersetSource):
     """

datahub/ingestion/source/qlik_sense/data_classes.py

@@ -1,8 +1,9 @@
+from copy import deepcopy
 from datetime import datetime
 from enum import Enum
 from typing import Dict, List, Optional, Type, Union
 
-from pydantic import BaseModel, Field, root_validator
+from pydantic import BaseModel, ConfigDict, Field, root_validator
 
 from datahub.emitter.mcp_builder import ContainerKey
 from datahub.ingestion.source.qlik_sense.config import QLIK_DATETIME_FORMAT, Constant
@@ -78,7 +79,11 @@ PERSONAL_SPACE_DICT = {
 }
 
 
-class Space(BaseModel):
+class _QlikBaseModel(BaseModel):
+    model_config = ConfigDict(coerce_numbers_to_str=True)
+
+
+class Space(_QlikBaseModel):
     id: str
     name: str
     description: str
@@ -89,6 +94,9 @@ class Space(BaseModel):
 
     @root_validator(pre=True)
     def update_values(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
+
         values[Constant.CREATEDAT] = datetime.strptime(
             values[Constant.CREATEDAT], QLIK_DATETIME_FORMAT
         )
@@ -98,7 +106,7 @@ class Space(BaseModel):
         return values
 
 
-class Item(BaseModel):
+class Item(_QlikBaseModel):
     id: str
     description: str = ""
     ownerId: str
@@ -107,7 +115,7 @@ class Item(BaseModel):
     updatedAt: datetime
 
 
-class SchemaField(BaseModel):
+class SchemaField(_QlikBaseModel):
     name: str
     dataType: Optional[str] = None
     primaryKey: Optional[bool] = None
@@ -115,6 +123,8 @@ class SchemaField(BaseModel):
 
     @root_validator(pre=True)
     def update_values(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         values[Constant.DATATYPE] = values.get(Constant.DATATYPE, {}).get(Constant.TYPE)
         return values
 
@@ -130,6 +140,8 @@ class QlikDataset(Item):
 
     @root_validator(pre=True)
     def update_values(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         # Update str time to datetime
         values[Constant.CREATEDAT] = datetime.strptime(
             values[Constant.CREATEDTIME], QLIK_DATETIME_FORMAT
@@ -148,13 +160,13 @@ class QlikDataset(Item):
         return values
 
 
-class AxisProperty(BaseModel):
+class AxisProperty(_QlikBaseModel):
     Title: str = Field(alias="qFallbackTitle")
     Min: str = Field(alias="qMin")
     Max: str = Field(alias="qMax")
 
 
-class Chart(BaseModel):
+class Chart(_QlikBaseModel):
     qId: str
     visualization: str
     title: str
@@ -164,13 +176,15 @@ class Chart(BaseModel):
 
     @root_validator(pre=True)
     def update_values(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         values[Constant.QID] = values[Constant.QINFO][Constant.QID]
         values["qDimension"] = values[Constant.HYPERCUBE]["qDimensionInfo"]
         values["qMeasure"] = values[Constant.HYPERCUBE]["qMeasureInfo"]
         return values
 
 
-class Sheet(BaseModel):
+class Sheet(_QlikBaseModel):
     id: str
     title: str
     description: str
@@ -181,6 +195,8 @@ class Sheet(BaseModel):
 
     @root_validator(pre=True)
     def update_values(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         values[Constant.CREATEDAT] = datetime.strptime(
             values[Constant.CREATEDDATE], QLIK_DATETIME_FORMAT
         )
@@ -190,7 +206,7 @@ class Sheet(BaseModel):
         return values
 
 
-class QlikTable(BaseModel):
+class QlikTable(_QlikBaseModel):
     tableName: str
     type: BoxType = Field(alias="boxType")
     tableAlias: str
@@ -206,6 +222,8 @@ class QlikTable(BaseModel):
 
     @root_validator(pre=True)
     def update_values(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         values[Constant.DATACONNECTORID] = values[Constant.CONNECTIONINFO][Constant.ID]
         values[Constant.DATACONNECTORPLATFORM] = values[Constant.CONNECTIONINFO][
             Constant.SOURCECONNECTORID
@@ -223,6 +241,8 @@ class App(Item):
 
     @root_validator(pre=True)
     def update_values(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
         values[Constant.CREATEDAT] = datetime.strptime(
             values[Constant.CREATEDDATE], QLIK_DATETIME_FORMAT
         )

datahub/ingestion/source/qlik_sense/qlik_sense.py

@@ -101,7 +101,7 @@ logger = logging.getLogger(__name__)
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
-    "Disabled by default. ",
+    "Disabled by default.",
 )
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(
@@ -109,6 +109,7 @@ logger = logging.getLogger(__name__)
     "Enabled by default, configured using `ingest_owner`",
 )
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class QlikSenseSource(StatefulIngestionSourceBase, TestableSource):
     """
     This plugin extracts the following:

datahub/ingestion/source/redash.py

@@ -447,7 +447,7 @@ class RedashSource(StatefulIngestionSourceBase):
                 dataset_urns = sql_parser_in_tables.in_tables
                 if sql_parser_in_tables.debug_info.table_error:
                     self.report.queries_problem_parsing.add(str(query_id))
-                    self.error(
+                    self.warn(
                         logger,
                         "sql-parsing",
                         f"exception {sql_parser_in_tables.debug_info.table_error} in parsing query-{query_id}-datasource-{data_source_id}",

datahub/ingestion/source/redshift/config.py

@@ -1,4 +1,5 @@
 import logging
+from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
@@ -6,9 +7,10 @@ from pydantic import root_validator
 from pydantic.fields import Field
 
 from datahub.configuration import ConfigModel
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.source_common import DatasetLineageProviderConfigBase
 from datahub.configuration.validate_field_removal import pydantic_removed_field
+from datahub.configuration.validate_field_rename import pydantic_renamed_field
 from datahub.ingestion.api.incremental_lineage_helper import (
     IncrementalLineageConfigMixin,
 )
@@ -94,13 +96,18 @@ class RedshiftConfig(
     # Because of this behavior, it uses dramatically fewer round trips for
     # large Redshift warehouses. As an example, see this query for the columns:
     # https://github.com/sqlalchemy-redshift/sqlalchemy-redshift/blob/60b4db04c1d26071c291aeea52f1dcb5dd8b0eb0/sqlalchemy_redshift/dialect.py#L745.
-    scheme: str = Field(
+    scheme: HiddenFromDocs[str] = Field(
         default="redshift+redshift_connector",
         description="",
-        hidden_from_docs=True,
     )
 
     _database_alias_removed = pydantic_removed_field("database_alias")
+    _use_lineage_v2_removed = pydantic_removed_field("use_lineage_v2")
+    _rename_lineage_v2_generate_queries_to_lineage_generate_queries = (
+        pydantic_renamed_field(
+            "lineage_v2_generate_queries", "lineage_generate_queries"
+        )
+    )
 
     default_schema: str = Field(
         default="public",
@@ -112,13 +119,9 @@ class RedshiftConfig(
         description="Whether target Redshift instance is serverless (alternative is provisioned cluster)",
     )
 
-    use_lineage_v2: bool = Field(
-        default=True,
-        description="Whether to use the new SQL-based lineage collector.",
-    )
-    lineage_v2_generate_queries: bool = Field(
+    lineage_generate_queries: bool = Field(
         default=True,
-        description="Whether to generate queries entities for the new SQL-based lineage collector.",
+        description="Whether to generate queries entities for the SQL-based lineage collector.",
     )
 
     include_table_lineage: bool = Field(
@@ -213,6 +216,9 @@ class RedshiftConfig(
 
     @root_validator(skip_on_failure=True)
     def connection_config_compatibility_set(cls, values: Dict) -> Dict:
+        # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
+        values = deepcopy(values)
+
         if (
             ("options" in values and "connect_args" in values["options"])
             and "extra_client_options" in values

datahub/ingestion/source/redshift/datashares.py

@@ -26,7 +26,7 @@ from datahub.utilities.search_utils import LogicalOperator
 
 class OutboundSharePlatformResource(BaseModel):
     namespace: str
-    platform_instance: Optional[str]
+    platform_instance: Optional[str] = None
     env: str
     source_database: str
     share_name: str