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

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

@@ -47,9 +47,11 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.entityLineage import ColumnLineage
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardServiceSource
@@ -88,7 +90,12 @@ class TableauSource(DashboardServiceSource):
     client: TableauClient
 
     @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: TableauConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, TableauConnection):
@@ -184,7 +191,7 @@ class TableauSource(DashboardServiceSource):
                     data_model_request = CreateDashboardDataModelRequest(
                         name=data_model.id,
                         displayName=data_model_name,
-                        service=self.context.dashboard_service,
+                        service=self.context.get().dashboard_service,
                         dataModelType=DataModelType.TableauDataModel.value,
                         serviceType=DashboardServiceType.Tableau.value,
                         columns=self.get_column_info(data_model),
@@ -228,19 +235,19 @@ class TableauSource(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 []
                 ],
                 dataModels=[
                     fqn.build(
                         self.metadata,
                         entity_type=DashboardDataModel,
-                        service_name=self.context.dashboard_service,
+                        service_name=self.context.get().dashboard_service,
                         data_model_name=data_model,
                     )
-                    for data_model in self.context.dataModels or []
+                    for data_model in self.context.get().dataModels or []
                 ],
                 tags=get_tag_labels(
                     metadata=self.metadata,
@@ -249,7 +256,7 @@ class TableauSource(DashboardServiceSource):
                     include_tags=self.source_config.includeTags,
                 ),
                 sourceUrl=dashboard_url,
-                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)
@@ -263,6 +270,50 @@ class TableauSource(DashboardServiceSource):
                 )
             )
 
+    @staticmethod
+    def _get_data_model_column_fqn(
+        data_model_entity: DashboardDataModel, column: str
+    ) -> Optional[str]:
+        """
+        Get fqn of column if exist in table entity
+        """
+        if not data_model_entity:
+            return None
+        for tbl_column in data_model_entity.columns:
+            for child_column in tbl_column.children or []:
+                if column.lower() == child_column.name.__root__.lower():
+                    return child_column.fullyQualifiedName.__root__
+        return None
+
+    def _get_column_lineage(
+        self,
+        upstream_table: UpstreamTable,
+        table_entity: Table,
+        data_model_entity: DashboardDataModel,
+        upstream_col_set: Set[str],
+    ) -> List[ColumnLineage]:
+        """
+        Get the column lineage from the fields
+        """
+        try:
+            column_lineage = []
+            for column in upstream_table.columns or []:
+                if column.id in upstream_col_set:
+                    from_column = get_column_fqn(
+                        table_entity=table_entity, column=column.name
+                    )
+                    to_column = self._get_data_model_column_fqn(
+                        data_model_entity=data_model_entity,
+                        column=column.id,
+                    )
+                    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}")
+            logger.debug(traceback.format_exc())
+
     def yield_dashboard_lineage_details(
         self, dashboard_details: TableauDashboard, db_service_name: str
     ) -> Iterable[Either[AddLineageRequest]]:
@@ -285,12 +336,22 @@ class TableauSource(DashboardServiceSource):
         for datamodel in dashboard_details.dataModels or []:
             try:
                 data_model_entity = self._get_datamodel(datamodel=datamodel)
+                upstream_col_set = {
+                    column.id
+                    for field in datamodel.fields
+                    for column in field.upstreamColumns
+                }
                 if data_model_entity:
                     for table in datamodel.upstreamTables or []:
                         om_table = self._get_database_table(db_service_entity, table)
                         if om_table:
+                            column_lineage = self._get_column_lineage(
+                                table, om_table, data_model_entity, upstream_col_set
+                            )
                             yield self._get_add_lineage_request(
-                                to_entity=data_model_entity, from_entity=om_table
+                                to_entity=data_model_entity,
+                                from_entity=om_table,
+                                column_lineage=column_lineage,
                             )
             except Exception as err:
                 yield Either(
@@ -340,7 +401,7 @@ class TableauSource(DashboardServiceSource):
                         classification_name=TABLEAU_TAG_CATEGORY,
                         include_tags=self.source_config.includeTags,
                     ),
-                    service=self.context.dashboard_service,
+                    service=self.context.get().dashboard_service,
                 )
                 yield Either(right=chart)
             except Exception as exc:
@@ -408,7 +469,7 @@ class TableauSource(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/database/athena/client.py

@@ -74,7 +74,7 @@ class AthenaLakeFormationClient:
             return LFTags(**response)
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(
+            logger.debug(
                 f"Unable to get lf tags for table resource [{table_name}] due to: {exc}"
             )
         return LFTags()

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

@@ -15,6 +15,7 @@ from typing import Iterable, Optional
 
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.source.database.athena.query_parser import (
+    QUERY_SUCCESS_STATUS,
     AthenaQueryParserSource,
 )
 from metadata.ingestion.source.database.lineage_source import LineageSource
@@ -37,6 +38,7 @@ class AthenaLineageSource(AthenaQueryParserSource, LineageSource):
                 if (
                     query.Status.SubmissionDateTime.date() >= self.start.date()
                     and self.is_not_dbt_or_om_query(query.Query)
+                    and query.Status.State.upper() == QUERY_SUCCESS_STATUS
                 ):
                     yield TableQuery(
                         query=query.Query,

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

@@ -12,7 +12,8 @@
 """Athena source module"""
 
 import traceback
-from typing import Iterable, Tuple
+from copy import deepcopy
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from pyathena.sqlalchemy.base import AthenaDialect
 from sqlalchemy import types
@@ -22,7 +23,8 @@ from sqlalchemy.engine.reflection import Inspector
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import (
     Column,
-    IntervalType,
+    PartitionColumnDetails,
+    PartitionIntervalTypes,
     Table,
     TablePartition,
     TableType,
@@ -57,6 +59,18 @@ logger = ingestion_logger()
 ATHENA_TAG = "ATHENA TAG"
 ATHENA_TAG_CLASSIFICATION = "ATHENA TAG CLASSIFICATION"
 
+ATHENA_INTERVAL_TYPE_MAP = {
+    **dict.fromkeys(["enum", "string", "VARCHAR"], PartitionIntervalTypes.COLUMN_VALUE),
+    **dict.fromkeys(
+        ["integer", "bigint", "INTEGER", "BIGINT"], PartitionIntervalTypes.INTEGER_RANGE
+    ),
+    **dict.fromkeys(
+        ["date", "timestamp", "DATE", "DATETIME", "TIMESTAMP"],
+        PartitionIntervalTypes.TIME_UNIT,
+    ),
+    "injected": PartitionIntervalTypes.INJECTED,
+}
+
 
 def _get_column_type(self, type_):
     """
@@ -123,6 +137,29 @@ def _get_column_type(self, type_):
     return col_type(*args)
 
 
+def _get_projection_details(
+    columns: List[Dict], projection_parameters: Dict
+) -> List[Dict]:
+    """Get the projection details for the columns
+
+    Args:
+        columns (List[Dict]): list of columns
+        projection_parameters (Dict): projection parameters
+    """
+    if not projection_parameters:
+        return columns
+
+    columns = deepcopy(columns)
+    for col in columns:
+        projection_details = next(
+            ({k: v} for k, v in projection_parameters.items() if k == col["name"]), None
+        )
+        if projection_details:
+            col["projection_type"] = projection_details[col["name"]]
+
+    return columns
+
+
 @reflection.cache
 def get_columns(self, connection, table_name, schema=None, **kw):
     """
@@ -147,6 +184,14 @@ def get_columns(self, connection, table_name, schema=None, **kw):
     ]
 
     if kw.get("only_partition_columns"):
+        # Return projected partition information to set partition type in `get_table_partition_details`
+        # projected partition fields are stored in the form of `projection.<field_name>.type` as a table parameter
+        projection_parameters = {
+            key_.split(".")[1]: value_
+            for key_, value_ in metadata.parameters.items()
+            if key_.startswith("projection") and key_.endswith("type")
+        }
+        columns = _get_projection_details(columns, projection_parameters)
         return columns
 
     columns += [
@@ -192,7 +237,9 @@ class AthenaSource(CommonDbSourceService):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: AthenaConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, AthenaConnection):
@@ -223,14 +270,34 @@ class AthenaSource(CommonDbSourceService):
 
     def get_table_partition_details(
         self, table_name: str, schema_name: str, inspector: Inspector
-    ) -> Tuple[bool, TablePartition]:
+    ) -> Tuple[bool, Optional[TablePartition]]:
+        """Get Athena table partition detail
+
+        Args:
+            table_name (str): name of the table
+            schema_name (str): name of the schema
+            inspector (Inspector):
+
+
+        Returns:
+            Tuple[bool, Optional[TablePartition]]:
+        """
         columns = inspector.get_columns(
             table_name=table_name, schema=schema_name, only_partition_columns=True
         )
         if columns:
             partition_details = TablePartition(
-                intervalType=IntervalType.COLUMN_VALUE.value,
-                columns=[column["name"] for column in columns],
+                columns=[
+                    PartitionColumnDetails(
+                        columnName=col["name"],
+                        intervalType=ATHENA_INTERVAL_TYPE_MAP.get(
+                            col.get("projection_type", str(col["type"])),
+                            PartitionIntervalTypes.COLUMN_VALUE,
+                        ),
+                        interval=None,
+                    )
+                    for col in columns
+                ]
             )
             return True, partition_details
         return False, None
@@ -251,8 +318,8 @@ class AthenaSource(CommonDbSourceService):
                         tag_fqn=fqn.build(
                             self.metadata,
                             DatabaseSchema,
-                            service_name=self.context.database_service,
-                            database_name=self.context.database,
+                            service_name=self.context.get().database_service,
+                            database_name=self.context.get().database,
                             schema_name=schema_name,
                         ),
                         tags=tag.TagValues,
@@ -280,7 +347,8 @@ class AthenaSource(CommonDbSourceService):
                 table_name, _ = table_name_and_type
                 table_tags = (
                     self.athena_lake_formation_client.get_table_and_column_tags(
-                        schema_name=self.context.database_schema, table_name=table_name
+                        schema_name=self.context.get().database_schema,
+                        table_name=table_name,
                     )
                 )
 
@@ -290,9 +358,9 @@ class AthenaSource(CommonDbSourceService):
                         tag_fqn=fqn.build(
                             self.metadata,
                             Table,
-                            service_name=self.context.database_service,
-                            database_name=self.context.database,
-                            schema_name=self.context.database_schema,
+                            service_name=self.context.get().database_service,
+                            database_name=self.context.get().database,
+                            schema_name=self.context.get().database_schema,
                             table_name=table_name,
                         ),
                         tags=tag.TagValues,
@@ -308,9 +376,9 @@ class AthenaSource(CommonDbSourceService):
                             tag_fqn=fqn.build(
                                 self.metadata,
                                 Column,
-                                service_name=self.context.database_service,
-                                database_name=self.context.database,
-                                schema_name=self.context.database_schema,
+                                service_name=self.context.get().database_service,
+                                database_name=self.context.get().database,
+                                schema_name=self.context.get().database_schema,
                                 table_name=table_name,
                                 column_name=column.Name,
                             ),

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

@@ -23,7 +23,7 @@ class QueryExecutionIdsResponse(BaseModel):
 
 
 class Status(BaseModel):
-    State: Optional[str]
+    State: Optional[str] = "FAILED"  # Default value
     SubmissionDateTime: Optional[datetime]
 
 

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

@@ -15,6 +15,7 @@ Athena Query parser module
 import traceback
 from abc import ABC
 from math import ceil
+from typing import Optional
 
 from metadata.clients.aws_client import AWSClient
 from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
@@ -40,6 +41,8 @@ ATHENA_QUERY_PAGINATOR_LIMIT = 50
 
 ATHENA_ENABLED_WORK_GROUP_STATE = "ENABLED"
 
+QUERY_SUCCESS_STATUS = "SUCCEEDED"
+
 
 class AthenaQueryParserSource(QueryParserSource, ABC):
     """
@@ -53,7 +56,9 @@ class AthenaQueryParserSource(QueryParserSource, ABC):
         self.client = AWSClient(self.service_connection.awsConfig).get_athena_client()
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: AthenaConnection = config.serviceConnection.__root__.config

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

@@ -15,6 +15,7 @@ from typing import Iterable
 
 from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
 from metadata.ingestion.source.database.athena.query_parser import (
+    QUERY_SUCCESS_STATUS,
     AthenaQueryParserSource,
 )
 from metadata.ingestion.source.database.usage_source import UsageSource
@@ -56,6 +57,7 @@ class AthenaUsageSource(AthenaQueryParserSource, UsageSource):
                 for query in query_list.QueryExecutions
                 if query.Status
                 and query.Query
+                and query.Status.State.upper() == QUERY_SUCCESS_STATUS
                 and query.Status.SubmissionDateTime.date() >= self.start.date()
                 and self.is_not_dbt_or_om_query(query.Query)
             ]

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

@@ -41,7 +41,10 @@ def get_connection_url(connection: Union[AzureSQLConnection, MssqlConnection]) -
     Build the connection URL
     """
 
-    if connection.authenticationMode:
+    if (
+        hasattr(connection.authenticationMode, "authentication")
+        and connection.authenticationMode.authentication is not None
+    ):
         connection_string = f"Driver={connection.driver};Server={connection.hostPort};Database={connection.database};"
         connection_string += f"Uid={connection.username};"
         if (

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

@@ -59,7 +59,9 @@ class AzuresqlSource(CommonDbSourceService, MultiDBSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: AzureSQLConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, AzureSQLConnection):
@@ -86,7 +88,7 @@ class AzuresqlSource(CommonDbSourceService, MultiDBSource):
                 database_fqn = fqn.build(
                     self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
+                    service_name=self.context.get().database_service,
                     database_name=new_database,
                 )
 

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

@@ -12,6 +12,7 @@
 AzureSQL usage module
 """
 from abc import ABC
+from typing import Optional
 
 from metadata.generated.schema.entity.services.connections.database.azureSQLConnection import (
     AzureSQLConnection,
@@ -32,7 +33,9 @@ class AzuresqlQueryParserSource(QueryParserSource, ABC):
     filters: str
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: AzureSQLConnection = config.serviceConnection.__root__.config

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

@@ -35,6 +35,9 @@ from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
+FK_CACHE = {}
+PK_CACHE = {}
+
 
 class InspectorWrapper(BaseModel):
     client: Any
@@ -80,14 +83,18 @@ def get_pk_constraint(
     This function overrides to get primary key constraint
     """
     try:
-        table_constraints = connection.engine.execute(
-            BIGQUERY_TABLE_CONSTRAINTS.format(
-                project_id=connection.engine.url.host,
-                schema_name=schema,
-                table_name=table_name,
+        constraints = PK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        if constraints is None:
+            constraints = connection.engine.execute(
+                BIGQUERY_TABLE_CONSTRAINTS.format(
+                    project_id=connection.engine.url.host,
+                    schema_name=schema,
+                )
             )
-        )
+            PK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+
         col_name = []
+        table_constraints = [row for row in constraints if row.table_name == table_name]
         for table_constraint in table_constraints:
             col_name.append(table_constraint.column_name)
         return {"constrained_columns": tuple(col_name)}
@@ -106,14 +113,18 @@ def get_foreign_keys(
     This function overrides to get foreign key constraint
     """
     try:
-        table_constraints = connection.engine.execute(
-            BIGQUERY_FOREIGN_CONSTRAINTS.format(
-                project_id=connection.engine.url.host,
-                schema_name=schema,
-                table_name=table_name,
+        constraints = FK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        if constraints is None:
+            constraints = connection.engine.execute(
+                BIGQUERY_FOREIGN_CONSTRAINTS.format(
+                    project_id=connection.engine.url.host,
+                    schema_name=schema,
+                )
             )
-        )
+            FK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+
         col_name = []
+        table_constraints = [row for row in constraints if row.table_name == table_name]
         for table_constraint in table_constraints:
             col_name.append(
                 {

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

@@ -0,0 +1,86 @@
+#  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.
+"""
+Bigquery Incremental Table processing logic
+"""
+from datetime import datetime
+from typing import List
+
+import google.cloud.logging
+from google.cloud.logging_v2.entries import LogEntry
+
+from metadata.ingestion.source.database.bigquery.models import (
+    BigQueryTable,
+    BigQueryTableMap,
+    SchemaName,
+    TableName,
+)
+from metadata.ingestion.source.database.bigquery.queries import (
+    BIGQUERY_GET_CHANGED_TABLES_FROM_CLOUD_LOGGING,
+)
+
+
+class BigQueryIncrementalTableProcessor:
+    def __init__(self, client: google.cloud.logging.Client):
+        self._client = client
+        self._changed_tables_map = BigQueryTableMap(table_map={})
+
+    @classmethod
+    def from_project(cls, project: str) -> "BigQueryIncrementalTableProcessor":
+        client = google.cloud.logging.Client(project=project)
+        return cls(client)
+
+    def _is_table_deleted(self, entry: LogEntry) -> bool:
+        if "tableDeletion" in entry.payload.get("metadata").keys():
+            return True
+        return False
+
+    def set_changed_tables_map(
+        self,
+        project: str,
+        dataset: str,
+        start_date: datetime,
+    ):
+        table_map = {}
+
+        resource_names = [f"projects/{project}"]
+        filters = BIGQUERY_GET_CHANGED_TABLES_FROM_CLOUD_LOGGING.format(
+            project=project,
+            dataset=dataset,
+            start_date=start_date.strftime("%Y-%m-%dT%H:%M:%SZ"),
+        )
+
+        entries = self._client.list_entries(
+            resource_names=resource_names,
+            filter_=filters,
+            order_by=google.cloud.logging.DESCENDING,
+        )
+
+        for entry in entries:
+            table_name = entry.payload.get("resourceName", "").split("/")[-1]
+            timestamp = entry.timestamp
+            deleted = self._is_table_deleted(entry)
+
+            if table_name not in table_map:
+                table_map[table_name] = BigQueryTable(
+                    name=table_name, timestamp=timestamp, deleted=deleted
+                )
+        self._changed_tables_map.add(dataset, table_map)
+
+    def get_deleted(self, schema_name: SchemaName) -> List[TableName]:
+        if self._changed_tables_map:
+            return self._changed_tables_map.get_deleted(schema_name)
+        return []
+
+    def get_not_deleted(self, schema_name: SchemaName) -> List[TableName]:
+        if self._changed_tables_map:
+            return self._changed_tables_map.get_not_deleted(schema_name)
+        return []