openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0__py3-none-any.whl

metadata/ingestion/source/dashboard/quicksight/metadata.py

@@ -11,7 +11,7 @@
 """QuickSight source module"""
 
 import traceback
-from typing import Any, Iterable, List, Optional
+from typing import Iterable, List, Optional
 
 from pydantic import ValidationError
 
@@ -34,7 +34,12 @@ 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.quicksight.models import DataSourceResp
+from metadata.ingestion.source.dashboard.quicksight.models import (
+    DashboardDetail,
+    DashboardResp,
+    DataSourceResp,
+    DescribeDataSourceResponse,
+)
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_chart
 from metadata.utils.logger import ingestion_logger
@@ -66,7 +71,9 @@ class QuicksightSource(DashboardServiceSource):
         }
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config = WorkflowSource.parse_obj(config_dict)
         connection: QuickSightConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, QuickSightConnection):
@@ -107,91 +114,95 @@ class QuicksightSource(DashboardServiceSource):
             dashboard["DashboardId"] for dashboard in dashboard_summary_list
         }
         dashboards = [
-            self.client.describe_dashboard(
-                AwsAccountId=self.aws_account_id, DashboardId=dashboard_id
-            )["Dashboard"]
+            DashboardResp(
+                **self.client.describe_dashboard(
+                    AwsAccountId=self.aws_account_id, DashboardId=dashboard_id
+                )
+            ).Dashboard
             for dashboard_id in dashboard_set
         ]
         return dashboards
 
-    def get_dashboard_name(self, dashboard: dict) -> str:
+    def get_dashboard_name(self, dashboard: DashboardDetail) -> str:
         """
         Get Dashboard Name
         """
-        return dashboard["Name"]
+        return dashboard.Name
 
-    def get_dashboard_details(self, dashboard: dict) -> dict:
+    def get_dashboard_details(self, dashboard: dict) -> DashboardDetail:
         """
         Get Dashboard Details
         """
-        return dashboard
+        return DashboardDetail(**dashboard)
 
     def yield_dashboard(
-        self, dashboard_details: dict
+        self, dashboard_details: DashboardDetail
     ) -> Iterable[Either[CreateDashboardRequest]]:
         """
         Method to Get Dashboard Entity
         """
         dashboard_request = CreateDashboardRequest(
-            name=dashboard_details["DashboardId"],
+            name=dashboard_details.DashboardId,
             sourceUrl=self.dashboard_url,
-            displayName=dashboard_details["Name"],
-            description=dashboard_details["Version"].get("Description"),
+            displayName=dashboard_details.Name,
+            description=dashboard_details.Version.Description
+            if dashboard_details.Version
+            else None,
             charts=[
                 fqn.build(
                     self.metadata,
                     entity_type=Chart,
-                    service_name=self.context.dashboard_service,
+                    service_name=self.context.get().dashboard_service,
                     chart_name=chart,
                 )
-                for chart in self.context.charts or []
+                for chart in self.context.get().charts or []
             ],
-            service=self.context.dashboard_service,
+            service=self.context.get().dashboard_service,
             owner=self.get_owner_ref(dashboard_details=dashboard_details),
         )
         yield Either(right=dashboard_request)
         self.register_record(dashboard_request=dashboard_request)
 
     def yield_dashboard_chart(
-        self, dashboard_details: Any
+        self, dashboard_details: DashboardDetail
     ) -> Iterable[Either[CreateChartRequest]]:
         """Get chart method"""
         # Each dashboard is guaranteed to have at least one sheet, which represents
         # a chart in the context of QuickSight
         # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/quicksight.html#QuickSight.Client.describe_dashboard
-        charts = dashboard_details["Version"]["Sheets"]
-        for chart in charts:
-            try:
-                if filter_by_chart(
-                    self.source_config.chartFilterPattern, chart["Name"]
-                ):
-                    self.status.filter(chart["Name"], "Chart Pattern not allowed")
-                    continue
+        if dashboard_details.Version:
+            for chart in dashboard_details.Version.Charts or []:
+                try:
+                    if filter_by_chart(
+                        self.source_config.chartFilterPattern, chart.Name
+                    ):
+                        self.status.filter(chart.Name, "Chart Pattern not allowed")
+                        continue
 
-                self.dashboard_url = (
-                    f"https://{self.aws_region}.quicksight.aws.amazon.com/sn/dashboards"
-                    f'/{dashboard_details.get("DashboardId")}'
-                )
-                yield Either(
-                    right=CreateChartRequest(
-                        name=chart["SheetId"],
-                        displayName=chart["Name"],
-                        chartType=ChartType.Other.value,
-                        sourceUrl=self.dashboard_url,
-                        service=self.context.dashboard_service,
+                    self.dashboard_url = (
+                        f"https://{self.aws_region}.quicksight.aws.amazon.com/sn/dashboards"
+                        f"/{dashboard_details.DashboardId}"
                     )
-                )
-            except Exception as exc:
-                yield Either(
-                    left=StackTraceError(
-                        name="Chart",
-                        error=f"Error creating chart [{chart}]: {exc}",
-                        stackTrace=traceback.format_exc(),
+                    yield Either(
+                        right=CreateChartRequest(
+                            name=chart.ChartId,
+                            displayName=chart.Name,
+                            chartType=ChartType.Other.value,
+                            sourceUrl=self.dashboard_url,
+                            service=self.context.get().dashboard_service,
+                        )
+                    )
+                except Exception as exc:
+                    yield Either(
+                        left=StackTraceError(
+                            name="Chart",
+                            error=f"Error creating chart [{chart}]: {exc}",
+                            stackTrace=traceback.format_exc(),
+                        )
                     )
-                )
 
     def yield_dashboard_lineage_details(  # pylint: disable=too-many-locals
-        self, dashboard_details: dict, db_service_name: str
+        self, dashboard_details: DashboardDetail, db_service_name: str
     ) -> Iterable[Either[AddLineageRequest]]:
         """
         Get lineage between dashboard and data sources
@@ -206,28 +217,26 @@ class QuicksightSource(DashboardServiceSource):
             )
             dataset_ids = {
                 dataset["DataSetId"]
-                for dataset in data_set_summary_list
-                if dataset.get("Arn") in dashboard_details["Version"]["DataSetArns"]
+                for dataset in data_set_summary_list or []
+                if dataset.get("Arn") in dashboard_details.Version.DataSetArns
             }
 
-            for dataset_id in dataset_ids:
-                for data_source in list(
-                    self.client.describe_data_set(
-                        AwsAccountId=self.aws_account_id, DataSetId=dataset_id
-                    )["DataSet"]["PhysicalTableMap"].values()
+            for dataset_id in dataset_ids or []:
+                for data_source in (
+                    list(
+                        self.client.describe_data_set(
+                            AwsAccountId=self.aws_account_id, DataSetId=dataset_id
+                        )["DataSet"]["PhysicalTableMap"].values()
+                    )
+                    or []
                 ):
                     try:
                         if not data_source.get("RelationalTable"):
                             raise KeyError(
                                 f"We currently don't support lineage to {list(data_source.keys())}"
                             )
-                        data_source_relational_table = data_source["RelationalTable"]
                         data_source_resp = DataSourceResp(
-                            datasource_arn=data_source_relational_table[
-                                "DataSourceArn"
-                            ],
-                            schema_name=data_source_relational_table["Schema"],
-                            table_name=data_source_relational_table["Name"],
+                            **data_source["RelationalTable"]
                         )
                     except (KeyError, ValidationError) as err:
                         yield Either(
@@ -255,42 +264,47 @@ class QuicksightSource(DashboardServiceSource):
 
                     data_source_ids = [
                         data_source_arn["DataSourceId"]
-                        for data_source_arn in data_source_summary_list
+                        for data_source_arn in data_source_summary_list or []
                         if data_source_arn["Arn"] in data_source_resp.datasource_arn
                     ]
 
-                    for data_source_id in data_source_ids:
-                        data_source_dict = self.client.describe_data_source(
-                            AwsAccountId=self.aws_account_id,
-                            DataSourceId=data_source_id,
-                        )["DataSource"]["DataSourceParameters"]
-                        for db in data_source_dict.keys():
-                            from_fqn = fqn.build(
-                                self.metadata,
-                                entity_type=Table,
-                                service_name=db_service_name,
-                                database_name=data_source_dict[db]["Database"],
-                                schema_name=schema_name,
-                                table_name=table_name,
-                                skip_es_search=True,
-                            )
-                            from_entity = self.metadata.get_by_name(
-                                entity=Table,
-                                fqn=from_fqn,
-                            )
-                            to_fqn = fqn.build(
-                                self.metadata,
-                                entity_type=Dashboard,
-                                service_name=self.config.serviceName,
-                                dashboard_name=dashboard_details["DashboardId"],
-                            )
-                            to_entity = self.metadata.get_by_name(
-                                entity=Dashboard,
-                                fqn=to_fqn,
-                            )
-                            yield self._get_add_lineage_request(
-                                to_entity=to_entity, from_entity=from_entity
+                    for data_source_id in data_source_ids or []:
+                        data_source_resp = DescribeDataSourceResponse(
+                            **self.client.describe_data_source(
+                                AwsAccountId=self.aws_account_id,
+                                DataSourceId=data_source_id,
                             )
+                        ).DataSource
+                        if data_source_resp and data_source_resp.DataSourceParameters:
+                            data_source_dict = data_source_resp.DataSourceParameters
+                            for db in data_source_dict.keys() or []:
+                                from_fqn = fqn.build(
+                                    self.metadata,
+                                    entity_type=Table,
+                                    service_name=db_service_name,
+                                    database_name=data_source_dict[db].get("Database"),
+                                    schema_name=schema_name,
+                                    table_name=table_name,
+                                    skip_es_search=True,
+                                )
+                                from_entity = self.metadata.get_by_name(
+                                    entity=Table,
+                                    fqn=from_fqn,
+                                )
+                                to_fqn = fqn.build(
+                                    self.metadata,
+                                    entity_type=Dashboard,
+                                    service_name=self.config.serviceName,
+                                    dashboard_name=dashboard_details.DashboardId,
+                                )
+                                to_entity = self.metadata.get_by_name(
+                                    entity=Dashboard,
+                                    fqn=to_fqn,
+                                )
+                                if from_entity is not None and to_entity is not None:
+                                    yield self._get_add_lineage_request(
+                                        to_entity=to_entity, from_entity=from_entity
+                                    )
         except Exception as exc:  # pylint: disable=broad-except
             yield Either(
                 left=StackTraceError(

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

@@ -12,11 +12,50 @@
 Pydantic Model to validate Quick Sight responses
 """
 
+from typing import List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class DataSourceResp(BaseModel):
-    datasource_arn: str
-    schema_name: str
-    table_name: str
+    datasource_arn: str = Field(alias="DataSourceArn")
+    schema_name: str = Field(alias="Schema")
+    table_name: str = Field(alias="Name")
+
+
+class VersionSheet(BaseModel):
+    ChartId: Optional[str] = Field(alias="SheetId")
+    Name: Optional[str]
+
+
+class DashboardVersion(BaseModel):
+    Status: Optional[str]
+    Arn: Optional[str]
+    SourceEntityArn: Optional[str]
+    DataSetArns: Optional[List]
+    Description: Optional[str]
+    Charts: Optional[List[VersionSheet]] = Field(alias="Sheets")
+
+
+class DashboardDetail(BaseModel):
+    DashboardId: str
+    Arn: Optional[str]
+    Name: str
+    Version: Optional[DashboardVersion]
+
+
+class DashboardResp(BaseModel):
+    Dashboard: DashboardDetail
+    Status: Optional[int]
+    RequestId: Optional[str]
+
+
+class DataSource(BaseModel):
+    DataSourceId: str
+    DataSourceParameters: Optional[dict]
+
+
+class DescribeDataSourceResponse(BaseModel):
+    DataSource: Optional[DataSource]
+    RequestId: Optional[str]
+    Status: Optional[int]

metadata/ingestion/source/dashboard/redash/metadata.py

@@ -68,7 +68,12 @@ class RedashSource(DashboardServiceSource):
         self.tags = []  # To create the tags before yielding final entities
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata):
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: RedashConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, RedashConnection):
@@ -154,12 +159,12 @@ class RedashSource(DashboardServiceSource):
                     fqn.build(
                         self.metadata,
                         entity_type=Chart,
-                        service_name=self.context.dashboard_service,
+                        service_name=self.context.get().dashboard_service,
                         chart_name=chart,
                     )
-                    for chart in self.context.charts or []
+                    for chart in self.context.get().charts or []
                 ],
-                service=self.context.dashboard_service,
+                service=self.context.get().dashboard_service,
                 sourceUrl=self.get_dashboard_url(dashboard_details),
                 tags=get_tag_labels(
                     metadata=self.metadata,
@@ -266,7 +271,7 @@ class RedashSource(DashboardServiceSource):
                         chartType=get_standard_chart_type(
                             visualization["type"] if visualization else ""
                         ),
-                        service=self.context.dashboard_service,
+                        service=self.context.get().dashboard_service,
                         sourceUrl=self.get_dashboard_url(dashboard_details),
                         description=visualization["description"]
                         if visualization

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

@@ -77,6 +77,11 @@ class SupersetAPISource(SupersetSourceMixin):
             dashboards = self.client.fetch_dashboards(current_page, page_size)
             current_page += 1
             for dashboard in dashboards.result:
+                if (
+                    not self.source_config.includeDraftDashboard
+                    and not dashboard.published
+                ):
+                    continue
                 yield dashboard
 
     def yield_dashboard(
@@ -94,12 +99,12 @@ class SupersetAPISource(SupersetSourceMixin):
                     fqn.build(
                         self.metadata,
                         entity_type=Chart,
-                        service_name=self.context.dashboard_service,
+                        service_name=self.context.get().dashboard_service,
                         chart_name=chart,
                     )
-                    for chart in self.context.charts or []
+                    for chart in self.context.get().charts or []
                 ],
-                service=self.context.dashboard_service,
+                service=self.context.get().dashboard_service,
                 owner=self.get_owner_ref(dashboard_details=dashboard_details),
             )
             yield Either(right=dashboard_request)
@@ -140,7 +145,7 @@ class SupersetAPISource(SupersetSourceMixin):
                     description=chart_json.description,
                     chartType=get_standard_chart_type(chart_json.viz_type),
                     sourceUrl=f"{clean_uri(self.service_connection.hostPort)}{chart_json.url}",
-                    service=self.context.dashboard_service,
+                    service=self.context.get().dashboard_service,
                 )
                 yield Either(right=chart)
             except Exception as exc:  # pylint: disable=broad-except
@@ -215,7 +220,7 @@ class SupersetAPISource(SupersetSourceMixin):
                     data_model_request = CreateDashboardDataModelRequest(
                         name=datasource_json.id,
                         displayName=datasource_json.result.table_name,
-                        service=self.context.dashboard_service,
+                        service=self.context.get().dashboard_service,
                         columns=self.get_column_info(datasource_json.result.columns),
                         dataModelType=DataModelType.SupersetDataModel.value,
                     )

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

@@ -45,6 +45,7 @@ from metadata.ingestion.source.dashboard.superset.queries import (
     FETCH_ALL_CHARTS,
     FETCH_COLUMN,
     FETCH_DASHBOARDS,
+    FETCH_PUBLISHED_DASHBOARDS,
 )
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_datamodel
@@ -100,7 +101,12 @@ class SupersetDBSource(SupersetSourceMixin):
         """
         Get List of all dashboards
         """
-        dashboards = self.engine.execute(FETCH_DASHBOARDS)
+        query = (
+            FETCH_DASHBOARDS
+            if self.source_config.includeDraftDashboard
+            else FETCH_PUBLISHED_DASHBOARDS
+        )
+        dashboards = self.engine.execute(query)
         for dashboard in dashboards:
             yield FetchDashboard(**dashboard)
 
@@ -117,12 +123,12 @@ class SupersetDBSource(SupersetSourceMixin):
                     fqn.build(
                         self.metadata,
                         entity_type=Chart,
-                        service_name=self.context.dashboard_service,
+                        service_name=self.context.get().dashboard_service,
                         chart_name=chart,
                     )
-                    for chart in self.context.charts or []
+                    for chart in self.context.get().charts or []
                 ],
-                service=self.context.dashboard_service,
+                service=self.context.get().dashboard_service,
                 owner=self.get_owner_ref(dashboard_details=dashboard_details),
             )
             yield Either(right=dashboard_request)
@@ -169,7 +175,7 @@ class SupersetDBSource(SupersetSourceMixin):
                     description=chart_json.description,
                     chartType=get_standard_chart_type(chart_json.viz_type),
                     sourceUrl=f"{clean_uri(self.service_connection.hostPort)}/explore/?slice_id={chart_json.id}",
-                    service=self.context.dashboard_service,
+                    service=self.context.get().dashboard_service,
                 )
                 yield Either(right=chart)
             except Exception as exc:
@@ -234,7 +240,7 @@ class SupersetDBSource(SupersetSourceMixin):
                     data_model_request = CreateDashboardDataModelRequest(
                         name=chart_json.datasource_id,
                         displayName=chart_json.table_name,
-                        service=self.context.dashboard_service,
+                        service=self.context.get().dashboard_service,
                         columns=self.get_column_info(col_names),
                         dataModelType=DataModelType.SupersetDataModel.value,
                     )

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

@@ -11,6 +11,8 @@
 """
 Superset source module
 """
+from typing import Optional
+
 from metadata.generated.schema.entity.services.connections.dashboard.supersetConnection import (
     SupersetConnection,
 )
@@ -32,7 +34,12 @@ class SupersetSource:
     """
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata):
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ):
         config = WorkflowSource.parse_obj(config_dict)
         connection: SupersetConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SupersetConnection):

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

@@ -70,7 +70,12 @@ class SupersetSourceMixin(DashboardServiceSource):
         self.all_charts = {}
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata):
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ):
         config = WorkflowSource.parse_obj(config_dict)
         connection: SupersetConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SupersetConnection):
@@ -177,9 +182,21 @@ class SupersetSourceMixin(DashboardServiceSource):
                             fqn=datamodel_fqn,
                         )
 
+                        datasource_json = self.client.fetch_datasource(
+                            chart_json.datasource_id
+                        )
+                        datasource_columns = self.get_column_info(
+                            datasource_json.result.columns
+                        )
+                        columns_list = [col.displayName for col in datasource_columns]
+                        column_lineage = self._get_column_lineage(
+                            from_entity, to_entity, columns_list
+                        )
                         if from_entity and to_entity:
                             yield self._get_add_lineage_request(
-                                to_entity=to_entity, from_entity=from_entity
+                                to_entity=to_entity,
+                                from_entity=from_entity,
+                                column_lineage=column_lineage,
                             )
                     except Exception as exc:
                         yield Either(
@@ -202,7 +219,7 @@ class SupersetSourceMixin(DashboardServiceSource):
         datamodel_fqn = fqn.build(
             self.metadata,
             entity_type=DashboardDataModel,
-            service_name=self.context.dashboard_service,
+            service_name=self.context.get().dashboard_service,
             data_model_name=datamodel.id,
         )
         if datamodel_fqn:

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

@@ -42,6 +42,7 @@ class DashboardResult(BaseModel):
     position_json: Optional[str]
     id: Optional[int]
     email: Optional[str]
+    published: Optional[bool]
 
 
 class SupersetDashboardCount(BaseModel):
@@ -135,6 +136,7 @@ class FetchDashboard(BaseModel):
     id: Optional[int]
     dashboard_title: Optional[str]
     position_json: Optional[str]
+    published: Optional[bool]
     email: Optional[str]
 
 

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

@@ -18,8 +18,8 @@ select
 	s.id,
 	s.slice_name,
 	s.description,
-    s.datasource_id,
-    s.viz_type,
+    	s.datasource_id,
+    	s.viz_type,
 	t.table_name,
 	t.schema,
 	db.database_name,
@@ -37,6 +37,7 @@ select
 	d.id, 
 	d.dashboard_title, 
 	d.position_json,
+    	d.published,
 	au.email 
 from 
 	dashboards d
@@ -46,6 +47,22 @@ ON
 	d.created_by_fk = au.id
 """
 
+FETCH_PUBLISHED_DASHBOARDS = """
+select
+	d.id, 
+	d.dashboard_title, 
+	d.position_json,
+    	d.published,
+	au.email 
+from 
+	dashboards d
+LEFT JOIN
+	ab_user au
+ON
+	d.created_by_fk = au.id
+where 
+	d.published=true
+"""
 
 FETCH_ALL_CHARTS_TEST = """
 select 
@@ -74,10 +91,10 @@ LIMIT 1
 FETCH_COLUMN = """
 select 
 	tc.id, 
-    table_name ,
-    column_name, 
-    type,
-    tc.description 
+    	table_name ,
+    	column_name, 
+    	type,
+    	tc.description 
 from 
 	table_columns  tc  
 inner join