openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

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

@@ -0,0 +1,67 @@
+#  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.
+
+"""
+Source connection handler
+"""
+
+from typing import Optional
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.dashboard.sigmaConnection import (
+    SigmaConnection,
+)
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
+from metadata.ingestion.connections.test_connections import (
+    SourceConnectionException,
+    test_connection_steps,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.dashboard.sigma.client import SigmaApiClient
+from metadata.utils.constants import THREE_MIN
+
+
+def get_connection(connection: SigmaConnection) -> SigmaApiClient:
+    """
+    Create connection
+    """
+    try:
+        return SigmaApiClient(connection)
+    except Exception as exc:
+        msg = f"Unknown error connecting with {connection}: {exc}."
+        raise SourceConnectionException(msg) from exc
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    client: SigmaApiClient,
+    service_connection: SigmaConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+
+    test_fn = {"GetToken": client.get_auth_token, "GetWorkbooks": client.get_dashboards}
+
+    return test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
+    )

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/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,10 @@
+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.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=AzuresqlSource,
+    lineage_source_class=AzuresqlLineageSource,
+    usage_source_class=AzuresqlUsageSource,
+)

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

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

@@ -11,14 +11,26 @@
 """
 Handle big query lineage extraction
 """
-from metadata.ingestion.source.database.bigquery.queries import BIGQUERY_STATEMENT
+from typing import Dict, List
+
+from metadata.ingestion.source.database.bigquery.queries import (
+    BIGQUERY_GET_STORED_PROCEDURE_QUERIES,
+    BIGQUERY_STATEMENT,
+)
 from metadata.ingestion.source.database.bigquery.query_parser import (
     BigqueryQueryParserSource,
 )
 from metadata.ingestion.source.database.lineage_source import LineageSource
+from metadata.ingestion.source.database.stored_procedures_mixin import (
+    QueryByProcedure,
+    StoredProcedureLineageMixin,
+)
+from metadata.utils.helpers import get_start_and_end
 
 
-class BigqueryLineageSource(BigqueryQueryParserSource, LineageSource):
+class BigqueryLineageSource(
+    BigqueryQueryParserSource, StoredProcedureLineageMixin, LineageSource
+):
     """
     Implements the necessary methods to extract
     Database lineage from Bigquery Source
@@ -32,3 +44,18 @@ class BigqueryLineageSource(BigqueryQueryParserSource, LineageSource):
             OR (statement_type = "INSERT" and UPPER(query) like '%%INSERT%%INTO%%SELECT%%')
         )
     """
+
+    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+        """
+        Pick the stored procedure name from the context
+        and return the list of associated queries
+        """
+        start, _ = get_start_and_end(self.source_config.queryLogDuration)
+        query = BIGQUERY_GET_STORED_PROCEDURE_QUERIES.format(
+            start_date=start,
+            region=self.service_connection.usageLocation,
+        )
+        queries_dict = self.procedure_queries_dict(
+            query=query,
+        )
+        return queries_dict

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

@@ -77,7 +77,6 @@ from metadata.ingestion.source.database.bigquery.models import (
     BigQueryStoredProcedure,
 )
 from metadata.ingestion.source.database.bigquery.queries import (
-    BIGQUERY_GET_STORED_PROCEDURE_QUERIES,
     BIGQUERY_GET_STORED_PROCEDURES,
     BIGQUERY_LIFE_CYCLE_QUERY,
     BIGQUERY_SCHEMA_DESCRIPTION,
@@ -95,14 +94,9 @@ from metadata.ingestion.source.database.life_cycle_query_mixin import (
     LifeCycleQueryMixin,
 )
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
-from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
-    StoredProcedureMixin,
-)
 from metadata.utils import fqn
 from metadata.utils.credentials import GOOGLE_CREDENTIALS
 from metadata.utils.filters import filter_by_database, filter_by_schema
-from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
     get_all_table_ddls,
@@ -223,9 +217,7 @@ Inspector.get_all_table_ddls = get_all_table_ddls
 Inspector.get_table_ddl = get_table_ddl
 
 
-class BigquerySource(
-    LifeCycleQueryMixin, StoredProcedureMixin, CommonDbSourceService, MultiDBSource
-):
+class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
     """
     Implements the necessary methods to extract
     Database metadata from Bigquery Source
@@ -850,22 +842,6 @@ class BigquerySource(
                 )
             )
 
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
-        """
-        Pick the stored procedure name from the context
-        and return the list of associated queries
-        """
-        start, _ = get_start_and_end(self.source_config.queryLogDuration)
-        query = BIGQUERY_GET_STORED_PROCEDURE_QUERIES.format(
-            start_date=start,
-            region=self.service_connection.usageLocation,
-        )
-        queries_dict = self.procedure_queries_dict(
-            query=query,
-        )
-
-        return queries_dict
-
     def mark_tables_as_deleted(self):
         """
         Use the current inspector to mark tables as deleted