acryl-datahub 1.0.0.2rc4__py3-none-any.whl → 1.0.0.3__py3-none-any.whl

datahub/ingestion/source/powerbi/config.py

@@ -192,6 +192,11 @@ class SupportedDataPlatform(Enum):
         datahub_data_platform_name="mysql",
     )
 
+    ODBC = DataPlatformPair(
+        powerbi_data_platform_name="Odbc",
+        datahub_data_platform_name="odbc",
+    )
+
 
 @dataclass
 class PowerBiDashboardSourceReport(StaleEntityRemovalSourceReport):
@@ -341,6 +346,13 @@ class PowerBiDashboardSourceConfig(
         "For Google BigQuery the datasource's server is google bigquery project name. "
         "For Databricks Unity Catalog the datasource's server is workspace FQDN.",
     )
+    # ODBC DSN to platform mapping
+    dsn_to_platform_name: Dict[str, str] = pydantic.Field(
+        default={},
+        description="A mapping of ODBC DSN to DataHub data platform name. "
+        "For example with an ODBC connection string 'DSN=database' where the database type "
+        "is 'PostgreSQL' you would configure the mapping as 'database: postgres'.",
+    )
     # deprecated warning
     _dataset_type_mapping = pydantic_field_deprecated(
         "dataset_type_mapping",
@@ -501,7 +513,7 @@ class PowerBiDashboardSourceConfig(
     include_workspace_name_in_dataset_urn: bool = pydantic.Field(
         default=False,
         description="It is recommended to set this to true, as it helps prevent the overwriting of datasets."
-        "Read section #11560 at https://datahubproject.io/docs/how/updating-datahub/ before enabling this option."
+        "Read section #11560 at https://docs.datahub.com/docs/how/updating-datahub/ before enabling this option."
         "To maintain backward compatibility, this is set to False.",
     )
 

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

@@ -75,3 +75,4 @@ class FunctionName(Enum):
     AMAZON_REDSHIFT_DATA_ACCESS = "AmazonRedshift.Database"
     DATABRICK_MULTI_CLOUD_DATA_ACCESS = "DatabricksMultiCloud.Catalogs"
     MYSQL_DATA_ACCESS = "MySQL.Database"
+    ODBC_DATA_ACCESS = "Odbc.DataSource"

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

@@ -0,0 +1,185 @@
+import re
+from typing import Optional, Tuple, Union
+
+server_patterns = [
+    r"Server=([^:]+)[:][0-9]+/.*",
+    r"SERVER=\{([^}]*)\}",
+    r"SERVER=([^;]*)",
+    r"HOST=\{([^}]*)\}",
+    r"HOST=([^;]*)",
+    r"DATA SOURCE=\{([^}]*)\}",
+    r"DATA SOURCE=([^;]*)",
+    r"DSN=\{([^}]*)\}",
+    r"DSN=([^;]*)",
+    r"Server=([^;]*)",
+    r"S3OutputLocation=([^;]*)",
+    r"HTTPPath=([^;]*)",
+    r"Host=([^;]*)",
+]
+
+dsn_patterns = [
+    r"DSN\s*=\s*\"([^\"]+)\"",
+    r"DSN\s*=\s*\'([^\']+)\'",
+    r"DSN\s*=\s*([^;]+)",
+]
+
+platform_patterns = {
+    "mysql": r"mysql",
+    "postgres": r"post(gre(s|sql)?|gres)",
+    "mssql": r"(sql\s*server|mssql|sqlncli)",
+    "oracle": r"oracle",
+    "db2": r"db2",
+    "sqlite": r"sqlite",
+    "access": r"(access|\.mdb|\.accdb)",
+    "excel": r"(excel|\.xls)",
+    "firebird": r"firebird",
+    "informix": r"informix",
+    "sybase": r"sybase",
+    "teradata": r"teradata",
+    "hadoop": r"(hadoop|hive)",
+    "snowflake": r"snowflake",
+    "redshift": r"redshift",
+    "bigquery": r"bigquery",
+    "athena": r"(athena|aws\s*athena)",
+    "databricks": r"(databricks|spark)",
+}
+
+powerbi_platform_names = {
+    "mysql": "MySQL",
+    "postgres": "PostgreSQL",
+    "mssql": "SQL Server",
+    "oracle": "Oracle",
+    "db2": "IBM DB2",
+    "sqlite": "SQLite",
+    "access": "Microsoft Access",
+    "excel": "Microsoft Excel",
+    "firebird": "Firebird",
+    "informix": "IBM Informix",
+    "sybase": "SAP Sybase",
+    "teradata": "Teradata",
+    "hadoop": "Hadoop",
+    "snowflake": "Snowflake",
+    "redshift": "Amazon Redshift",
+    "bigquery": "Google BigQuery",
+    "athena": "Amazon Athena",
+    "databricks": "Databricks",
+}
+
+
+def extract_driver(connection_string: str) -> Union[str, None]:
+    """
+    Parse an ODBC connection string and extract the driver name.
+    Handles whitespace in driver names and various connection string formats.
+
+    Args:
+        connection_string (str): The ODBC connection string
+
+    Returns:
+        str: The extracted driver name, or None if not found
+    """
+    # Match DRIVER={driver name} pattern
+    driver_match = re.search(r"DRIVER=\{([^}]*)}", connection_string, re.IGNORECASE)
+
+    if driver_match:
+        return driver_match.group(1).strip()
+
+    # Alternative pattern for DRIVER=driver
+    driver_match = re.search(r"DRIVER=([^;]*)", connection_string, re.IGNORECASE)
+
+    if driver_match:
+        return driver_match.group(1).strip()
+
+    return None
+
+
+def extract_dsn(connection_string: str) -> Union[str, None]:
+    """
+    Extract the DSN value from an ODBC connection string.
+
+    Args:
+        connection_string (str): The ODBC connection string
+
+    Returns:
+        str or None: The extracted DSN value, or None if not found
+    """
+    for pattern in dsn_patterns:
+        match = re.search(pattern, connection_string, re.IGNORECASE)
+        if match:
+            return match.group(1).strip()
+
+    return None
+
+
+def extract_server(connection_string: str) -> Union[str, None]:
+    """
+    Parse an ODBC connection string and extract the server name.
+    Handles various parameter names for server (SERVER, Host, Data Source, etc.)
+
+    Args:
+        connection_string (str): The ODBC connection string
+
+    Returns:
+        str: The extracted server name, or None if not found
+    """
+    for pattern in server_patterns:
+        server_match = re.search(pattern, connection_string, re.IGNORECASE)
+        if server_match:
+            return server_match.group(1).strip()
+
+    # Special case for Athena: extract from AwsRegion if no server found
+    region_match = re.search(r"AwsRegion=([^;]*)", connection_string, re.IGNORECASE)
+    if region_match:
+        return f"aws-athena-{region_match.group(1).strip()}"
+
+    # Special case for Databricks: try to extract hostname from JDBC URL
+    jdbc_match = re.search(r"jdbc:spark://([^:;/]+)", connection_string, re.IGNORECASE)
+    if jdbc_match:
+        return jdbc_match.group(1).strip()
+
+    return None
+
+
+def extract_platform(connection_string: str) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Extract the database platform name from the ODBC driver name.
+    Returns the lowercase platform name.
+
+    Args:
+        connection_string (str): The ODBC connection string
+
+    Returns:
+        tuple: A tuple containing the normalized platform name and the corresponding
+        Power BI platform name, or None if not recognized.
+    """
+    driver_name = extract_driver(connection_string)
+    if not driver_name:
+        return None, None
+
+    driver_lower = driver_name.lower()
+
+    for platform, pattern in platform_patterns.items():
+        if re.search(pattern, driver_lower):
+            return platform, powerbi_platform_names.get(platform)
+
+    return None, None
+
+
+def normalize_platform_name(platform: str) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Normalizes the platform name by matching it with predefined patterns and maps it to
+    a corresponding Power BI platform name.
+
+    Args:
+        platform (str): The platform name to normalize
+
+    Returns:
+        tuple: A tuple containing the normalized platform name and the corresponding
+        Power BI platform name, or None if not recognized.
+    """
+    platform_lower = platform.lower()
+
+    for platform, pattern in platform_patterns.items():
+        if re.search(pattern, platform_lower):
+            return platform, powerbi_platform_names.get(platform)
+
+    return None, None

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

@@ -29,6 +29,12 @@ from datahub.ingestion.source.powerbi.m_query.data_classes import (
     Lineage,
     ReferencedTable,
 )
+from datahub.ingestion.source.powerbi.m_query.odbc import (
+    extract_dsn,
+    extract_platform,
+    extract_server,
+    normalize_platform_name,
+)
 from datahub.ingestion.source.powerbi.rest_api_wrapper.data_classes import Table
 from datahub.metadata.schema_classes import SchemaFieldDataTypeClass
 from datahub.sql_parsing.sqlglot_lineage import (
@@ -155,6 +161,7 @@ class AbstractLineage(ABC):
                 tree_function.token_values(arg_list)
             ),
         )
+        logger.debug(f"DB Details: {arguments}")
 
         if len(arguments) < 2:
             logger.debug(f"Expected minimum 2 arguments, but got {len(arguments)}")
@@ -940,6 +947,147 @@ class NativeQueryLineage(AbstractLineage):
         )
 
 
+class OdbcLineage(AbstractLineage):
+    def create_lineage(
+        self, data_access_func_detail: DataAccessFunctionDetail
+    ) -> Lineage:
+        logger.debug(
+            f"Processing {self.get_platform_pair().powerbi_data_platform_name} "
+            f"data-access function detail {data_access_func_detail}"
+        )
+
+        connect_string, _ = self.get_db_detail_from_argument(
+            data_access_func_detail.arg_list
+        )
+
+        if not connect_string:
+            self.reporter.warning(
+                title="Can not extract ODBC connect string",
+                message="Can not extract ODBC connect string from data access function. Skipping Lineage creation.",
+                context=f"table-name={self.table.full_name}, data-access-func-detail={data_access_func_detail}",
+            )
+            return Lineage.empty()
+
+        logger.debug(f"ODBC connect string: {connect_string}")
+        data_platform, powerbi_platform = extract_platform(connect_string)
+        server_name = extract_server(connect_string)
+
+        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:
+                logger.debug(f"Attempting to map DSN {dsn} to platform")
+                name = self.config.dsn_to_platform_name.get(dsn)
+                if name:
+                    logger.debug(f"Found DSN {dsn} mapped to platform {name}")
+                    data_platform, powerbi_platform = normalize_platform_name(name)
+
+        if not data_platform or not powerbi_platform:
+            self.reporter.warning(
+                title="Can not determine ODBC platform",
+                message="Can not determine platform from ODBC connect string. Skipping Lineage creation.",
+                context=f"table-name={self.table.full_name}, connect-string={connect_string}",
+            )
+            return Lineage.empty()
+
+        platform_pair: DataPlatformPair = self.create_platform_pair(
+            data_platform, powerbi_platform
+        )
+
+        if not server_name and self.config.server_to_platform_instance:
+            self.reporter.warning(
+                title="Can not determine ODBC server name",
+                message="Can not determine server name with server_to_platform_instance mapping. Skipping Lineage creation.",
+                context=f"table-name={self.table.full_name}",
+            )
+            return Lineage.empty()
+        elif not server_name:
+            server_name = "unknown"
+
+        database_name = None
+        schema_name = None
+        table_name = None
+        qualified_table_name = None
+
+        temp_accessor: Optional[IdentifierAccessor] = (
+            data_access_func_detail.identifier_accessor
+        )
+
+        while temp_accessor:
+            logger.debug(
+                f"identifier = {temp_accessor.identifier} items = {temp_accessor.items}"
+            )
+            if temp_accessor.items.get("Kind") == "Database":
+                database_name = temp_accessor.items["Name"]
+
+            if temp_accessor.items.get("Kind") == "Schema":
+                schema_name = temp_accessor.items["Name"]
+
+            if temp_accessor.items.get("Kind") == "Table":
+                table_name = temp_accessor.items["Name"]
+
+            if temp_accessor.next is not None:
+                temp_accessor = temp_accessor.next
+            else:
+                break
+
+        if (
+            database_name is not None
+            and schema_name is not None
+            and table_name is not None
+        ):
+            qualified_table_name = f"{database_name}.{schema_name}.{table_name}"
+        elif database_name is not None and table_name is not None:
+            qualified_table_name = f"{database_name}.{table_name}"
+
+        if not qualified_table_name:
+            self.reporter.warning(
+                title="Can not determine qualified table name",
+                message="Can not determine qualified table name for ODBC data source. Skipping Lineage creation.",
+                context=f"table-name={self.table.full_name}, data-platform={data_platform}",
+            )
+            logger.warning(
+                f"Can not determine qualified table name for ODBC data source {data_platform} "
+                f"table {self.table.full_name}."
+            )
+            return Lineage.empty()
+
+        logger.debug(
+            f"ODBC Platform {data_platform} found qualified table name {qualified_table_name}"
+        )
+
+        urn = make_urn(
+            config=self.config,
+            platform_instance_resolver=self.platform_instance_resolver,
+            data_platform_pair=platform_pair,
+            server=server_name,
+            qualified_table_name=qualified_table_name,
+        )
+
+        column_lineage = self.create_table_column_lineage(urn)
+
+        return Lineage(
+            upstreams=[
+                DataPlatformTable(
+                    data_platform_pair=platform_pair,
+                    urn=urn,
+                )
+            ],
+            column_lineage=column_lineage,
+        )
+
+    @staticmethod
+    def create_platform_pair(
+        data_platform: str, powerbi_platform: str
+    ) -> DataPlatformPair:
+        return DataPlatformPair(data_platform, powerbi_platform)
+
+    def get_platform_pair(self) -> DataPlatformPair:
+        return SupportedDataPlatform.ODBC.value
+
+
 class SupportedPattern(Enum):
     DATABRICKS_QUERY = (
         DatabricksLineage,
@@ -991,6 +1139,11 @@ class SupportedPattern(Enum):
         FunctionName.NATIVE_QUERY,
     )
 
+    ODBC = (
+        OdbcLineage,
+        FunctionName.ODBC_DATA_ACCESS,
+    )
+
     def handler(self) -> Type[AbstractLineage]:
         return self.value[0]
 

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

@@ -63,10 +63,10 @@ class SessionWithTimeout(requests.Session):
         super().__init__(*args, **kwargs)
         self.timeout = timeout
 
-    def request(self, method, url, **kwargs):
+    def request(self, method, url, *args, **kwargs):
         # Set the default timeout if none is provided
         kwargs.setdefault("timeout", self.timeout)
-        return super().request(method, url, **kwargs)
+        return super().request(method, url, *args, **kwargs)
 
 
 class DataResolverBase(ABC):

datahub/ingestion/source/redshift/usage.py

@@ -182,15 +182,16 @@ class RedshiftUsageExtractor:
         self.report.num_operational_stats_filtered = 0
 
         if self.config.include_operational_stats:
-            with self.report.new_stage(USAGE_EXTRACTION_OPERATIONAL_STATS):
-                with PerfTimer() as timer:
-                    # Generate operation aspect workunits
-                    yield from self._gen_operation_aspect_workunits(
-                        self.connection, all_tables
-                    )
-                    self.report.operational_metadata_extraction_sec[
-                        self.config.database
-                    ] = timer.elapsed_seconds(digits=2)
+            with self.report.new_stage(
+                USAGE_EXTRACTION_OPERATIONAL_STATS
+            ), PerfTimer() as timer:
+                # Generate operation aspect workunits
+                yield from self._gen_operation_aspect_workunits(
+                    self.connection, all_tables
+                )
+                self.report.operational_metadata_extraction_sec[
+                    self.config.database
+                ] = timer.elapsed_seconds(digits=2)
 
         # Generate aggregate events
         with self.report.new_stage(USAGE_EXTRACTION_USAGE_AGGREGATION):

datahub/ingestion/source/sigma/config.py

@@ -1,8 +1,9 @@
 import logging
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import pydantic
+from pydantic import BaseModel, Field
 
 from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.source_common import (
@@ -53,15 +54,82 @@ class Constant:
     DEFAULT_API_URL = "https://aws-api.sigmacomputing.com/v2"
 
 
+class WorkspaceCounts(BaseModel):
+    workbooks_count: int = 0
+    datasets_count: int = 0
+    elements_count: int = 0
+    pages_count: int = 0
+
+    def is_empty(self) -> bool:
+        return (
+            self.workbooks_count == 0
+            and self.datasets_count == 0
+            and self.elements_count == 0
+            and self.pages_count == 0
+        )
+
+    def as_obj(self) -> dict:
+        return {
+            "workbooks_count": self.workbooks_count,
+            "datasets_count": self.datasets_count,
+            "elements_count": self.elements_count,
+            "pages_count": self.pages_count,
+        }
+
+
+class SigmaWorkspaceEntityFilterReport(EntityFilterReport):
+    type: str = "workspace"
+
+    workspace_counts: Dict[str, WorkspaceCounts] = Field(
+        default_factory=dict,
+        description="Counts of workbooks, datasets, elements and pages in each workspace.",
+    )
+
+    def increment_workbooks_count(self, workspace_id: str) -> None:
+        if workspace_id not in self.workspace_counts:
+            self.workspace_counts[workspace_id] = WorkspaceCounts()
+        self.workspace_counts[workspace_id].workbooks_count += 1
+
+    def increment_datasets_count(self, workspace_id: str) -> None:
+        if workspace_id not in self.workspace_counts:
+            self.workspace_counts[workspace_id] = WorkspaceCounts()
+        self.workspace_counts[workspace_id].datasets_count += 1
+
+    def increment_elements_count(self, workspace_id: str) -> None:
+        if workspace_id not in self.workspace_counts:
+            self.workspace_counts[workspace_id] = WorkspaceCounts()
+        self.workspace_counts[workspace_id].elements_count += 1
+
+    def increment_pages_count(self, workspace_id: str) -> None:
+        if workspace_id not in self.workspace_counts:
+            self.workspace_counts[workspace_id] = WorkspaceCounts()
+        self.workspace_counts[workspace_id].pages_count += 1
+
+    def as_obj(self) -> dict:
+        return {
+            "filtered": self.dropped_entities.as_obj(),
+            "processed": self.processed_entities.as_obj(),
+            "workspace_counts": {
+                key: item.as_obj() for key, item in self.workspace_counts.items()
+            },
+        }
+
+
 @dataclass
 class SigmaSourceReport(StaleEntityRemovalSourceReport):
-    workspaces: EntityFilterReport = EntityFilterReport.field(type="workspace")
-    number_of_workspaces: Optional[int] = None
+    workspaces: SigmaWorkspaceEntityFilterReport = field(
+        default_factory=SigmaWorkspaceEntityFilterReport
+    )
     non_accessible_workspaces_count: int = 0
-    shared_entities_count: int = 0
-    number_of_datasets: int = 0
-    number_of_workbooks: int = 0
+
+    datasets: EntityFilterReport = EntityFilterReport.field(type="dataset")
+    datasets_without_workspace: int = 0
+
+    workbooks: EntityFilterReport = EntityFilterReport.field(type="workbook")
+    workbooks_without_workspace: int = 0
+
     number_of_files_metadata: Dict[str, int] = field(default_factory=dict)
+    empty_workspaces: List[str] = field(default_factory=list)
 
 
 class PlatformDetail(PlatformInstanceConfigMixin, EnvConfigMixin):

datahub/ingestion/source/sigma/sigma.py

@@ -35,6 +35,7 @@ from datahub.ingestion.source.sigma.config import (
     PlatformDetail,
     SigmaSourceConfig,
     SigmaSourceReport,
+    WorkspaceCounts,
 )
 from datahub.ingestion.source.sigma.data_classes import (
     Element,
@@ -163,7 +164,6 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
     def _get_allowed_workspaces(self) -> List[Workspace]:
         all_workspaces = self.sigma_api.workspaces.values()
         logger.info(f"Number of workspaces = {len(all_workspaces)}")
-        self.reporter.number_of_workspaces = len(all_workspaces)
 
         allowed_workspaces = []
         for workspace in all_workspaces:
@@ -285,6 +285,7 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
         yield self._gen_dataset_properties(dataset_urn, dataset)
 
         if dataset.workspaceId:
+            self.reporter.workspaces.increment_datasets_count(dataset.workspaceId)
             yield from add_entity_to_container(
                 container_key=self._gen_workspace_key(dataset.workspaceId),
                 entity_type="dataset",
@@ -468,6 +469,8 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
             ).as_workunit()
 
             if workbook.workspaceId:
+                self.reporter.workspaces.increment_elements_count(workbook.workspaceId)
+
                 yield self._gen_entity_browsepath_aspect(
                     entity_urn=chart_urn,
                     parent_entity_urn=builder.make_container_urn(
@@ -525,6 +528,7 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
             all_input_fields: List[InputFieldClass] = []
 
             if workbook.workspaceId:
+                self.reporter.workspaces.increment_pages_count(workbook.workspaceId)
                 yield self._gen_entity_browsepath_aspect(
                     entity_urn=dashboard_urn,
                     parent_entity_urn=builder.make_container_urn(
@@ -614,6 +618,8 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
 
         paths = workbook.path.split("/")[1:]
         if workbook.workspaceId:
+            self.reporter.workspaces.increment_workbooks_count(workbook.workspaceId)
+
             yield self._gen_entity_browsepath_aspect(
                 entity_urn=dashboard_urn,
                 parent_entity_urn=builder.make_container_urn(
@@ -667,6 +673,15 @@ class SigmaSource(StatefulIngestionSourceBase, TestableSource):
                 f"{workspace.name} ({workspace.workspaceId})"
             )
             yield from self._gen_workspace_workunit(workspace)
+            if self.reporter.workspaces.workspace_counts.get(
+                workspace.workspaceId, WorkspaceCounts()
+            ).is_empty():
+                logger.warning(
+                    f"Workspace {workspace.name} ({workspace.workspaceId}) is empty. If this is not expected, add the user associated with the Client ID/Secret to each workspace with missing metadata"
+                )
+                self.reporter.empty_workspaces.append(
+                    f"{workspace.name} ({workspace.workspaceId})"
+                )
         yield from self._gen_sigma_dataset_upstream_lineage_workunit()
 
     def get_report(self) -> SourceReport: