openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/ingestion/source/dashboard/sigma/metadata.py

@@ -0,0 +1,331 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""Sigma source module"""
+
+import traceback
+from typing import Iterable, List, Optional
+
+from metadata.generated.schema.api.data.createChart import CreateChartRequest
+from metadata.generated.schema.api.data.createDashboard import CreateDashboardRequest
+from metadata.generated.schema.api.data.createDashboardDataModel import (
+    CreateDashboardDataModelRequest,
+)
+from metadata.generated.schema.entity.data.chart import Chart
+from metadata.generated.schema.entity.data.dashboardDataModel import (
+    DashboardDataModel,
+    DataModelType,
+)
+from metadata.generated.schema.entity.data.table import Column, DataType, Table
+from metadata.generated.schema.entity.services.connections.dashboard.sigmaConnection import (
+    SigmaConnection,
+)
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    Markdown,
+    SourceUrl,
+)
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.dashboard.dashboard_service import DashboardServiceSource
+from metadata.ingestion.source.dashboard.sigma.models import (
+    Elements,
+    NodeDetails,
+    Workbook,
+    WorkbookDetails,
+)
+from metadata.utils import fqn
+from metadata.utils.filters import filter_by_chart
+from metadata.utils.helpers import get_standard_chart_type
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class SigmaSource(DashboardServiceSource):
+    """
+    Sigma Source Class
+    """
+
+    @classmethod
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ):
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SigmaConnection = config.serviceConnection.root.config
+        if not isinstance(connection, SigmaConnection):
+            raise InvalidSourceException(
+                f"Expected SigmaConnection, but got {connection}"
+            )
+        return cls(config, metadata)
+
+    def __init__(
+        self,
+        config: WorkflowSource,
+        metadata: OpenMetadata,
+    ):
+        super().__init__(config, metadata)
+        self.data_models: List[Elements] = []
+
+    def get_dashboards_list(self) -> Optional[List[Workbook]]:
+        """
+        get list of dashboard
+        """
+        return self.client.get_dashboards()
+
+    def get_dashboard_name(self, dashboard: Workbook) -> Optional[str]:
+        """
+        get dashboard name
+        """
+        return dashboard.name
+
+    def get_dashboard_details(self, dashboard: Workbook) -> Optional[WorkbookDetails]:
+        """
+        get dashboard details
+        """
+        return self.client.get_dashboard_detail(dashboard.workbookId)
+
+    def yield_dashboard(
+        self, dashboard_details: WorkbookDetails
+    ) -> Iterable[Either[CreateDashboardRequest]]:
+        """
+        yield Dashboard Entity
+        """
+        try:
+            dashboard_request = CreateDashboardRequest(
+                name=EntityName(str(dashboard_details.workbookId)),
+                displayName=dashboard_details.name,
+                description=Markdown(dashboard_details.description)
+                if dashboard_details.description
+                else None,
+                charts=[
+                    FullyQualifiedEntityName(
+                        fqn.build(
+                            self.metadata,
+                            entity_type=Chart,
+                            service_name=self.context.get().dashboard_service,
+                            chart_name=chart,
+                        )
+                    )
+                    for chart in self.context.get().charts or []
+                ],
+                service=FullyQualifiedEntityName(self.context.get().dashboard_service),
+                sourceUrl=SourceUrl(dashboard_details.url),
+                owners=self.get_owner_ref(dashboard_details=dashboard_details),
+            )
+            yield Either(right=dashboard_request)
+            self.register_record(dashboard_request=dashboard_request)
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name="Dashboard",
+                    error=f"Error to yield dashboard for {dashboard_details}: {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
+    def yield_dashboard_chart(
+        self, dashboard_details: WorkbookDetails
+    ) -> Iterable[Either[CreateChartRequest]]:
+        """
+        yield dashboard charts
+        """
+        charts = self.client.get_chart_details(dashboard_details.workbookId)
+        for chart in charts or []:
+            try:
+                if filter_by_chart(self.source_config.chartFilterPattern, chart.name):
+                    self.status.filter(chart.name, "Chart Pattern not allowed")
+                    continue
+                yield Either(
+                    right=CreateChartRequest(
+                        name=EntityName(str(chart.elementId)),
+                        displayName=chart.name,
+                        chartType=get_standard_chart_type(chart.vizualizationType),
+                        service=FullyQualifiedEntityName(
+                            self.context.get().dashboard_service
+                        ),
+                        sourceUrl=SourceUrl(dashboard_details.url),
+                        description=Markdown(dashboard_details.description)
+                        if dashboard_details.description
+                        else None,
+                    )
+                )
+            except Exception as exc:
+                yield Either(
+                    left=StackTraceError(
+                        name="Chart",
+                        error=(
+                            "Error to yield dashboard chart for : "
+                            f"{chart.elementId} and {dashboard_details}: {exc}"
+                        ),
+                        stackTrace=traceback.format_exc(),
+                    )
+                )
+
+    def _get_datamodel(self, datamodel_id: str):
+        datamodel_fqn = fqn.build(
+            self.metadata,
+            entity_type=DashboardDataModel,
+            service_name=self.context.get().dashboard_service,
+            data_model_name=datamodel_id,
+        )
+        if datamodel_fqn:
+            return self.metadata.get_by_name(
+                entity=DashboardDataModel,
+                fqn=datamodel_fqn,
+            )
+        return None
+
+    def _get_table_entity_from_node(
+        self, node: NodeDetails, db_service_name: str
+    ) -> Optional[Table]:
+        """
+        Get the table entity for lineage
+        """
+        if node.node_schema:
+            try:
+                table_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=Table,
+                    service_name=db_service_name,
+                    schema_name=node.node_schema,
+                    table_name=node.name,
+                    database_name="",
+                )
+                if table_fqn:
+                    return self.metadata.get_by_name(
+                        entity=Table,
+                        fqn=table_fqn,
+                    )
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.warning(f"Error occured while finding table fqn: {exc}")
+        return None
+
+    def yield_dashboard_lineage_details(
+        self, dashboard_details: WorkbookDetails, db_service_name: Optional[str]
+    ):
+        """
+        yield dashboard lineage
+        """
+        if not db_service_name:
+            return
+        # charts and datamodels are same here as we are using charts as metadata for datamodels
+        for data_model in self.data_models or []:
+            try:
+                data_model_entity = self._get_datamodel(
+                    datamodel_id=data_model.elementId
+                )
+                if data_model_entity:
+                    nodes = self.client.get_lineage_details(
+                        dashboard_details.workbookId, data_model.elementId
+                    )
+                    for node in nodes:
+                        table_entity = self._get_table_entity_from_node(
+                            node, db_service_name
+                        )
+                        if table_entity and data_model.columns:
+                            columns_list = data_model.columns
+                            column_lineage = self._get_column_lineage(
+                                table_entity, data_model_entity, columns_list
+                            )
+                            yield self._get_add_lineage_request(
+                                to_entity=data_model_entity,
+                                from_entity=table_entity,
+                                column_lineage=column_lineage,
+                            )
+            except Exception as exc:
+                yield Either(
+                    left=StackTraceError(
+                        name=f"{dashboard_details.name} Lineage",
+                        error=(
+                            "Error to yield dashboard lineage details for DB "
+                            f"service name [{db_service_name}]: {exc}"
+                        ),
+                        stackTrace=traceback.format_exc(),
+                    )
+                )
+
+    def get_column_info(self, element: Elements) -> Optional[List[Column]]:
+        """Build data model columns"""
+        datamodel_columns = []
+        for col in element.columns or []:
+            try:
+                datamodel_columns.append(
+                    Column(
+                        name=col,
+                        displayName=col,
+                        dataType=DataType.UNKNOWN,
+                        dataTypeDisplay="Sigma Field",
+                    )
+                )
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.warning(f"Error to yield datamodel column: {exc}")
+        return datamodel_columns
+
+    def yield_datamodel(
+        self, dashboard_details: WorkbookDetails
+    ) -> Iterable[Either[DashboardDataModel]]:
+        if self.source_config.includeDataModels:
+            # we are ingesting charts/Elements as datamodels here
+            self.data_models = self.client.get_chart_details(
+                dashboard_details.workbookId
+            )
+            for data_model in self.data_models or []:
+                try:
+                    data_model_request = CreateDashboardDataModelRequest(
+                        name=EntityName(data_model.elementId),
+                        displayName=data_model.name,
+                        service=FullyQualifiedEntityName(
+                            self.context.get().dashboard_service
+                        ),
+                        dataModelType=DataModelType.SigmaDataModel.value,
+                        serviceType=self.service_connection.type.value,
+                        columns=self.get_column_info(data_model),
+                    )
+                    yield Either(right=data_model_request)
+                    self.register_record_datamodel(datamodel_request=data_model_request)
+                except Exception as exc:
+                    yield Either(
+                        left=StackTraceError(
+                            name=data_model.elementId,
+                            error=f"Error yielding Data Model [{data_model.elementId}]: {exc}",
+                            stackTrace=traceback.format_exc(),
+                        )
+                    )
+
+    def get_owner_ref(
+        self, dashboard_details: WorkbookDetails
+    ) -> Optional[EntityReferenceList]:
+        """
+        Get owner from email
+        """
+        try:
+            if dashboard_details.ownerId:
+                owner = self.client.get_owner_detail(dashboard_details.ownerId)
+                return self.metadata.get_reference_by_email(owner.email)
+            return None
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not fetch owner data due to {err}")
+        return None

metadata/ingestion/source/dashboard/sigma/models.py

@@ -0,0 +1,97 @@
+#  Copyright 2023 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+PowerBI Models
+"""
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class AuthToken(BaseModel):
+    access_token: str
+    refresh_token: Optional[str] = None
+    token_type: str
+    expires_in: Optional[int] = 0
+
+
+class Workbook(BaseModel):
+    workbookId: str
+    name: Optional[str] = None
+    ownerId: Optional[str] = None
+
+
+class WorkbookDetails(BaseModel):
+    workbookId: str
+    name: Optional[str] = None
+    createdAt: str
+    url: str
+    path: Optional[str] = None
+    ownerId: Optional[str] = None
+    isArchived: bool
+    description: Optional[str] = None
+
+
+class WorkBookResponseDetails(BaseModel):
+    entries: Optional[List[Workbook]] = []
+
+
+class OwnerDetails(BaseModel):
+    organizationId: str
+    email: str
+
+
+class WorkBookPage(BaseModel):
+    pageId: str
+
+
+class WorkBookPageResponse(BaseModel):
+    entries: Optional[List[WorkBookPage]] = []
+
+
+class Elements(BaseModel):
+    elementId: str
+    name: Optional[str] = None
+    vizualizationType: Optional[str] = None
+    columns: Optional[List[str]] = []
+
+
+class ElementsResponse(BaseModel):
+    entries: Optional[List[Elements]] = []
+
+
+class EdgeSource(BaseModel):
+    source: str
+
+    @property
+    def node_id(self):
+        if self.source:
+            if "inode-" in self.source:
+                return self.source.replace("inode-", "")
+        return None
+
+
+class EdgeSourceResponse(BaseModel):
+    edges: Optional[List[EdgeSource]] = []
+
+
+class NodeDetails(BaseModel):
+    id: str
+    name: Optional[str]
+    node_type: str = Field(alias="type")
+    path: Optional[str] = ""
+
+    @property
+    def node_schema(self):
+        if self.node_type == "table" and self.path:
+            if "/" in self.path:
+                return self.path.split("/", maxsplit=-1)[-1]
+        return None

metadata/ingestion/source/dashboard/sigma/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.dashboard.sigma.metadata import SigmaSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=SigmaSource)

metadata/ingestion/source/dashboard/superset/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.dashboard.superset.metadata import SupersetSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=SupersetSource)

metadata/ingestion/source/dashboard/tableau/metadata.py

@@ -69,6 +69,7 @@ from metadata.ingestion.source.dashboard.tableau.models import (
     ChartUrl,
     DataSource,
     DatasourceField,
+    TableAndQuery,
     TableauDashboard,
     TableauTag,
     UpstreamTable,
@@ -149,7 +150,7 @@ class TableauSource(DashboardServiceSource):
         return None
 
     @staticmethod
-    def _get_data_models_tags(dataModels: [DataSource]) -> Set[str]:
+    def _get_data_models_tags(dataModels: List[DataSource]) -> Set[str]:
         """
         Get the tags from the data model in the upstreamDatasources
         """
@@ -382,9 +383,12 @@ class TableauSource(DashboardServiceSource):
                         column=column.id,
                     )
                     for to_column in to_columns:
-                        column_lineage.append(
-                            ColumnLineage(fromColumns=[from_column], toColumn=to_column)
-                        )
+                        if from_column and to_column:
+                            column_lineage.append(
+                                ColumnLineage(
+                                    fromColumns=[from_column], toColumn=to_column
+                                )
+                            )
             return column_lineage
         except Exception as exc:
             logger.debug(f"Error to get column lineage: {exc}")
@@ -451,17 +455,22 @@ class TableauSource(DashboardServiceSource):
                 column.id
                 for field in upstream_data_model.fields
                 for column in field.upstreamColumns
+                if column is not None
             }
             for table in datamodel.upstreamTables or []:
                 om_tables = self._get_database_tables(db_service_entity, table)
-                for om_table in om_tables or []:
+                for om_table_and_query in om_tables or []:
                     column_lineage = self._get_column_lineage(
-                        table, om_table, upstream_data_model_entity, upstream_col_set
+                        table,
+                        om_table_and_query.table,
+                        upstream_data_model_entity,
+                        upstream_col_set,
                     )
                     yield self._get_add_lineage_request(
                         to_entity=upstream_data_model_entity,
-                        from_entity=om_table,
+                        from_entity=om_table_and_query.table,
                         column_lineage=column_lineage,
+                        sql=om_table_and_query.query,
                     )
         except Exception as err:
             yield Either(
@@ -698,7 +707,7 @@ class TableauSource(DashboardServiceSource):
 
     def _get_table_entities_from_api(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         In case we get the table details from the Graphql APIs we process them
         """
@@ -734,7 +743,7 @@ class TableauSource(DashboardServiceSource):
                     fqn=table_fqn,
                 )
                 if table_entity:
-                    return [table_entity]
+                    return [TableAndQuery(table=table_entity)]
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Error to get tables for lineage using GraphQL Apis: {exc}")
@@ -742,7 +751,7 @@ class TableauSource(DashboardServiceSource):
 
     def _get_table_entities_from_query(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         In case we get the table details from the Graphql APIs we process them
         """
@@ -778,7 +787,12 @@ class TableauSource(DashboardServiceSource):
                         database_schema=schema_name,
                         table=table_name,
                     )
-                    tables_list.extend(from_entities)
+                    tables_list.extend(
+                        [
+                            TableAndQuery(table=table, query=custom_sql_table.query)
+                            for table in from_entities
+                        ]
+                    )
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
@@ -787,7 +801,7 @@ class TableauSource(DashboardServiceSource):
 
     def _get_database_tables(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         Get the table entities for lineage
         """
@@ -879,6 +893,9 @@ class TableauSource(DashboardServiceSource):
         try:
             return dashboard_details.project.name
         except Exception as exc:
+            logger.info(
+                f"Cannot parse project name for dashboard:{dashboard_details.id} from Tableau server"
+            )
             logger.debug(traceback.format_exc())
             logger.warning(
                 f"Error fetching project name for {dashboard_details.id}: {exc}"

metadata/ingestion/source/dashboard/tableau/models.py

@@ -18,6 +18,7 @@ from typing import Any, Dict, List, Optional, Union
 from pydantic import BaseModel, ConfigDict, Field, validator
 
 from metadata.generated.schema.entity.data.chart import ChartType
+from metadata.generated.schema.entity.data.table import Table
 
 
 class TableauBaseModel(BaseModel):
@@ -172,3 +173,12 @@ class TableauDashboard(TableauBaseModel):
     webpageUrl: Optional[str] = None
     charts: Optional[List[TableauChart]] = None
     dataModels: List[DataSource] = []
+
+
+class TableAndQuery(BaseModel):
+    """
+    Wrapper class for Table entity and associated Query for lineage
+    """
+
+    table: Table
+    query: Optional[str] = None

metadata/ingestion/source/dashboard/tableau/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.dashboard.tableau.metadata import TableauSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=TableauSource)

metadata/ingestion/source/database/athena/service_spec.py

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.athena.lineage import AthenaLineageSource
+from metadata.ingestion.source.database.athena.metadata import AthenaSource
+from metadata.ingestion.source.database.athena.usage import AthenaUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=AthenaSource,
+    lineage_source_class=AthenaLineageSource,
+    usage_source_class=AthenaUsageSource,
+)

metadata/ingestion/source/database/azuresql/service_spec.py

@@ -0,0 +1,12 @@
+from metadata.ingestion.source.database.azuresql.lineage import AzuresqlLineageSource
+from metadata.ingestion.source.database.azuresql.metadata import AzuresqlSource
+from metadata.ingestion.source.database.azuresql.usage import AzuresqlUsageSource
+from metadata.sampler.sqlalchemy.azuresql.sampler import AzureSQLSampler
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=AzuresqlSource,
+    lineage_source_class=AzuresqlLineageSource,
+    usage_source_class=AzuresqlUsageSource,
+    sampler_class=AzureSQLSampler,
+)

metadata/ingestion/source/database/bigquery/connection.py

@@ -17,6 +17,7 @@ from datetime import datetime
 from functools import partial
 from typing import Optional
 
+from google.api_core.exceptions import NotFound
 from google.cloud.datacatalog_v1 import PolicyTagManagerClient
 from sqlalchemy.engine import Engine
 
@@ -154,8 +155,8 @@ def test_connection(
         test_fn = {
             "CheckAccess": partial(test_connection_engine_step, engine),
             "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),
-            "GetTables": partial(execute_inspector_func, engine, "get_table_names"),
-            "GetViews": partial(execute_inspector_func, engine, "get_view_names"),
+            "GetTables": partial(get_table_view_names, engine),
+            "GetViews": partial(get_table_view_names, engine),
             "GetTags": test_tags,
             "GetQueries": partial(
                 test_query,
@@ -176,3 +177,28 @@ def test_connection(
         )
 
     return test_connection_inner(engine)
+
+
+def get_table_view_names(connection, schema=None):
+    with connection.connect() as conn:
+        current_schema = schema
+        client = conn.connection._client
+        item_types = ["TABLE", "EXTERNAL", "VIEW", "MATERIALIZED_VIEW"]
+        datasets = client.list_datasets()
+        result = []
+        for dataset in datasets:
+            if current_schema is not None and current_schema != dataset.dataset_id:
+                continue
+
+            try:
+                tables = client.list_tables(dataset.reference, page_size=1)
+                for table in tables:
+                    if table.table_type in item_types:
+                        break
+            except NotFound:
+                # It's possible that the dataset was deleted between when we
+                # fetched the list of datasets and when we try to list the
+                # tables from it. See:
+                # https://github.com/googleapis/python-bigquery-sqlalchemy/issues/105
+                pass
+        return result