openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

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

@@ -15,6 +15,10 @@ GraphQL queries used during ingestion
 
 TABLEAU_DATASOURCES_QUERY = """
 {{
+workbooks(filter:{{luid: "{workbook_id}"}}){{
+  id
+  luid
+  name
   embeddedDatasourcesConnection(first: {first}, offset: {offset} ) {{
     nodes {{
       id
@@ -29,11 +33,6 @@ TABLEAU_DATASOURCES_QUERY = """
         }}
         description
       }}
-      workbook {{
-        id
-        luid
-        name
-      }}
       upstreamTables {{
         id
         luid
@@ -57,5 +56,6 @@ TABLEAU_DATASOURCES_QUERY = """
     }}
     totalCount
   }}
+  }}
 }}
 """

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

@@ -0,0 +1,80 @@
+#  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.
+"""
+Wrapper module of Athena client
+"""
+import traceback
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
+    AthenaConnection,
+)
+from metadata.ingestion.models.lf_tags_model import LFTags
+from metadata.ingestion.source.database.athena.connection import (
+    get_lake_formation_client,
+)
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+
+class AthenaLakeFormationClient:
+    """
+    Athena Lake Formation Client
+    """
+
+    def __init__(
+        self,
+        connection: AthenaConnection,
+    ):
+        self.lake_formation_client = get_lake_formation_client(connection=connection)
+
+    def get_database_tags(self, name: str) -> Optional[LFTags]:
+        """
+        Method to call the API and get the database tags
+        """
+        try:
+            response = self.lake_formation_client.get_resource_lf_tags(
+                Resource={"Database": {"Name": name}}
+            )
+            lf_tags = LFTags(**response)
+            return lf_tags.LFTagOnDatabase
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unable to get lf tags for database resource [{name}] due to: {exc}"
+            )
+        return None
+
+    def get_table_and_column_tags(self, schema_name: str, table_name: str) -> LFTags:
+        """
+        Method to call the API and get the table and column tags
+        """
+        try:
+            response = self.lake_formation_client.get_resource_lf_tags(
+                Resource={
+                    "Table": {
+                        "DatabaseName": schema_name,
+                        "Name": table_name,
+                    },
+                    "TableWithColumns": {
+                        "DatabaseName": schema_name,
+                        "Name": table_name,
+                    },
+                }
+            )
+            return LFTags(**response)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unable to get lf tags for table resource [{table_name}] due to: {exc}"
+            )
+        return LFTags()

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

@@ -75,6 +75,13 @@ def get_connection(connection: AthenaConnection) -> Engine:
     )
 
 
+def get_lake_formation_client(connection: AthenaConnection):
+    """
+    Get the lake formation client
+    """
+    return AWSClient(connection.awsConfig).get_lake_formation_client()
+
+
 def test_connection(
     metadata: OpenMetadata,
     engine: Engine,

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

@@ -11,32 +11,52 @@
 
 """Athena source module"""
 
-from typing import Iterable
+import traceback
+from typing import Iterable, Tuple
 
 from pyathena.sqlalchemy.base import AthenaDialect
 from sqlalchemy import types
 from sqlalchemy.engine import reflection
+from sqlalchemy.engine.reflection import Inspector
 
-from metadata.generated.schema.entity.data.table import TableType
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    IntervalType,
+    Table,
+    TablePartition,
+    TableType,
+)
 from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
     AthenaConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source import sqa_types
+from metadata.ingestion.source.database.athena.client import AthenaLakeFormationClient
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import is_complex_type
+from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
 logger = ingestion_logger()
 
+ATHENA_TAG = "ATHENA TAG"
+ATHENA_TAG_CLASSIFICATION = "ATHENA TAG CLASSIFICATION"
+
 
 def _get_column_type(self, type_):
     """
@@ -73,20 +93,14 @@ def _get_column_type(self, type_):
         "struct": sqa_types.SQAStruct,
         "row": sqa_types.SQAStruct,
         "map": sqa_types.SQAMap,
+        "decimal": types.DECIMAL,
+        "varchar": types.VARCHAR,
+        "char": types.CHAR,
     }
-    if name in ["decimal"]:
-        col_type = types.DECIMAL
-        if length:
-            precision, scale = length.split(",")
-            args = [int(precision), int(scale)]
-    elif name in ["char"]:
-        col_type = types.CHAR
+    if name in ["decimal", "char", "varchar"]:
+        col_type = col_map[name]
         if length:
-            args = [int(length)]
-    elif name in ["varchar"]:
-        col_type = types.VARCHAR
-        if length:
-            args = [int(length)]
+            args = [int(l) for l in length.split(",")]
     elif type_.startswith("array"):
         parsed_type = (
             ColumnTypeParser._parse_datatype_string(  # pylint: disable=protected-access
@@ -94,7 +108,13 @@ def _get_column_type(self, type_):
             )
         )
         col_type = col_map["array"]
-        args = [col_map.get(parsed_type.get("arrayDataType").lower(), types.String)]
+        if parsed_type["arrayDataType"].lower().startswith("array"):
+            # as OpenMetadata doesn't store any details on children of array, we put
+            # in type as string as default to avoid Array item_type required issue
+            # from sqlalchemy types
+            args = [types.String]
+        else:
+            args = [col_map.get(parsed_type.get("arrayDataType").lower(), types.String)]
     elif col_map.get(name):
         col_type = col_map.get(name)
     else:
@@ -121,10 +141,14 @@ def get_columns(self, connection, table_name, schema=None, **kw):
             "comment": c.comment,
             "system_data_type": c.type,
             "is_complex": is_complex_type(c.type),
-            "dialect_options": {"awsathena_partition": None},
+            "dialect_options": {"awsathena_partition": True},
         }
-        for c in metadata.columns
+        for c in metadata.partition_keys
     ]
+
+    if kw.get("only_partition_columns"):
+        return columns
+
     columns += [
         {
             "name": c.name,
@@ -135,10 +159,11 @@ def get_columns(self, connection, table_name, schema=None, **kw):
             "comment": c.comment,
             "system_data_type": c.type,
             "is_complex": is_complex_type(c.type),
-            "dialect_options": {"awsathena_partition": True},
+            "dialect_options": {"awsathena_partition": None},
         }
-        for c in metadata.partition_keys
+        for c in metadata.columns
     ]
+
     return columns
 
 
@@ -176,6 +201,16 @@ class AthenaSource(CommonDbSourceService):
             )
         return cls(config, metadata)
 
+    def __init__(
+        self,
+        config: WorkflowSource,
+        metadata: OpenMetadata,
+    ):
+        super().__init__(config, metadata)
+        self.athena_lake_formation_client = AthenaLakeFormationClient(
+            connection=self.service_connection
+        )
+
     def query_table_names_and_types(
         self, schema_name: str
     ) -> Iterable[TableNameAndType]:
@@ -185,3 +220,110 @@ class AthenaSource(CommonDbSourceService):
             TableNameAndType(name=name, type_=TableType.External)
             for name in self.inspector.get_table_names(schema_name)
         ]
+
+    def get_table_partition_details(
+        self, table_name: str, schema_name: str, inspector: Inspector
+    ) -> Tuple[bool, 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],
+            )
+            return True, partition_details
+        return False, None
+
+    def yield_tag(
+        self, schema_name: str
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        Method to yield schema tags
+        """
+        if self.source_config.includeTags:
+            try:
+                tags = self.athena_lake_formation_client.get_database_tags(
+                    name=schema_name
+                )
+                for tag in tags or []:
+                    yield from get_ometa_tag_and_classification(
+                        tag_fqn=fqn.build(
+                            self.metadata,
+                            DatabaseSchema,
+                            service_name=self.context.database_service,
+                            database_name=self.context.database,
+                            schema_name=schema_name,
+                        ),
+                        tags=tag.TagValues,
+                        classification_name=tag.TagKey,
+                        tag_description=ATHENA_TAG,
+                        classification_description=ATHENA_TAG_CLASSIFICATION,
+                    )
+            except Exception as exc:
+                yield Either(
+                    left=StackTraceError(
+                        name="Tags and Classifications",
+                        error=f"Failed to fetch database tags due to [{exc}]",
+                        stackTrace=traceback.format_exc(),
+                    )
+                )
+
+    def yield_table_tags(
+        self, table_name_and_type: Tuple[str, TableType]
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        Method to yield table and column tags
+        """
+        if self.source_config.includeTags:
+            try:
+                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
+                    )
+                )
+
+                # yield the table tags
+                for tag in table_tags.LFTagsOnTable or []:
+                    yield from get_ometa_tag_and_classification(
+                        tag_fqn=fqn.build(
+                            self.metadata,
+                            Table,
+                            service_name=self.context.database_service,
+                            database_name=self.context.database,
+                            schema_name=self.context.database_schema,
+                            table_name=table_name,
+                        ),
+                        tags=tag.TagValues,
+                        classification_name=tag.TagKey,
+                        tag_description=ATHENA_TAG,
+                        classification_description=ATHENA_TAG_CLASSIFICATION,
+                    )
+
+                # yield the column tags
+                for column in table_tags.LFTagsOnColumns or []:
+                    for tag in column.LFTags or []:
+                        yield from get_ometa_tag_and_classification(
+                            tag_fqn=fqn.build(
+                                self.metadata,
+                                Column,
+                                service_name=self.context.database_service,
+                                database_name=self.context.database,
+                                schema_name=self.context.database_schema,
+                                table_name=table_name,
+                                column_name=column.Name,
+                            ),
+                            tags=tag.TagValues,
+                            classification_name=tag.TagKey,
+                            tag_description=ATHENA_TAG,
+                            classification_description=ATHENA_TAG_CLASSIFICATION,
+                        )
+            except Exception as exc:
+                yield Either(
+                    left=StackTraceError(
+                        name="Tags and Classifications",
+                        error=f"Failed to fetch table/column tags due to [{exc}]",
+                        stackTrace=traceback.format_exc(),
+                    )
+                )

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

@@ -15,12 +15,13 @@ Source connection handler
 from typing import Optional, Union
 from urllib.parse import quote_plus
 
-from sqlalchemy.engine import Engine
+from sqlalchemy.engine import URL, Engine
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
 from metadata.generated.schema.entity.services.connections.database.azureSQLConnection import (
+    Authentication,
     AzureSQLConnection,
 )
 from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
@@ -40,13 +41,29 @@ def get_connection_url(connection: Union[AzureSQLConnection, MssqlConnection]) -
     Build the connection URL
     """
 
+    if connection.authenticationMode:
+        connection_string = f"Driver={connection.driver};Server={connection.hostPort};Database={connection.database};"
+        connection_string += f"Uid={connection.username};"
+        if (
+            connection.authenticationMode.authentication
+            == Authentication.ActiveDirectoryPassword
+        ):
+            connection_string += f"Pwd={connection.password.get_secret_value()};"
+
+        connection_string += f"Encrypt={'yes' if connection.authenticationMode.encrypt else 'no'};TrustServerCertificate={'yes' if connection.authenticationMode.trustServerCertificate else 'no'};"
+        connection_string += f"Connection Timeout={connection.authenticationMode.connectionTimeout or 30};Authentication={connection.authenticationMode.authentication.value};"
+
+        connection_url = URL.create(
+            "mssql+pyodbc", query={"odbc_connect": connection_string}
+        )
+        return connection_url
     url = f"{connection.scheme.value}://"
 
     if connection.username:
         url += f"{quote_plus(connection.username)}"
         url += (
             f":{quote_plus(connection.password.get_secret_value())}"
-            if connection
+            if connection.password
             else ""
         )
         url += "@"
@@ -54,12 +71,13 @@ def get_connection_url(connection: Union[AzureSQLConnection, MssqlConnection]) -
     url += f"{connection.hostPort}"
     url += f"/{quote_plus(connection.database)}" if connection.database else ""
     url += f"?driver={quote_plus(connection.driver)}"
+
     options = get_connection_options_dict(connection)
     if options:
         if not connection.database:
             url += "/"
         params = "&".join(
-            f"{key}={quote_plus(value)}" for (key, value) in options.items() if value
+            f"{key}={quote_plus(value)}" for key, value in options.items() if value
         )
         url = f"{url}?{params}"
 

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

@@ -77,7 +77,6 @@ class AzuresqlSource(CommonDbSourceService, MultiDBSource):
         yield from self._execute_database_query(AZURE_SQL_GET_DATABASES)
 
     def get_database_names(self) -> Iterable[str]:
-
         if not self.config.serviceConnection.__root__.config.ingestAllDatabases:
             configured_db = self.config.serviceConnection.__root__.config.database
             self.set_inspector(database_name=configured_db)

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

@@ -44,6 +44,9 @@ from metadata.ingestion.connections.test_connections import (
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.bigquery.queries import BIGQUERY_TEST_STATEMENT
 from metadata.utils.credentials import set_google_credentials
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 def get_connection_url(connection: BigQueryConnection) -> str:
@@ -107,9 +110,27 @@ def test_connection(
             return policy_tags
 
     def test_tags():
-        taxonomies = PolicyTagManagerClient().list_taxonomies(
-            parent=f"projects/{engine.url.host}/locations/{service_connection.taxonomyLocation}"
-        )
+        taxonomy_project_ids = []
+        if engine.url.host:
+            taxonomy_project_ids.append(engine.url.host)
+        if service_connection.taxonomyProjectID:
+            taxonomy_project_ids.extend(service_connection.taxonomyProjectID)
+        if not taxonomy_project_ids:
+            logger.info("'taxonomyProjectID' is not set, so skipping this test.")
+            return None
+
+        taxonomy_location = service_connection.taxonomyLocation
+        if not taxonomy_location:
+            logger.info("'taxonomyLocation' is not set, so skipping this test.")
+            return None
+
+        taxonomies = []
+        for project_id in taxonomy_project_ids:
+            taxonomies.extend(
+                PolicyTagManagerClient().list_taxonomies(
+                    parent=f"projects/{project_id}/locations/{taxonomy_location}"
+                )
+            )
         return get_tags(taxonomies)
 
     def test_connection_inner(engine):

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

@@ -12,7 +12,7 @@
 """
 Source connection helper
 """
-
+import traceback
 from typing import Any
 
 from pydantic import BaseModel
@@ -26,7 +26,14 @@ from metadata.generated.schema.security.credentials.gcpValues import (
     SingleProjectId,
 )
 from metadata.ingestion.source.connections import get_connection
+from metadata.ingestion.source.database.bigquery.queries import (
+    BIGQUERY_FOREIGN_CONSTRAINTS,
+    BIGQUERY_TABLE_CONSTRAINTS,
+)
 from metadata.utils.bigquery_utils import get_bigquery_client
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class InspectorWrapper(BaseModel):
@@ -64,3 +71,63 @@ def get_inspector_details(
     inspector = inspect(engine)
 
     return InspectorWrapper(client=client, engine=engine, inspector=inspector)
+
+
+def get_pk_constraint(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    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,
+            )
+        )
+        col_name = []
+        for table_constraint in table_constraints:
+            col_name.append(table_constraint.column_name)
+        return {"constrained_columns": tuple(col_name)}
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(
+            f"Error while fetching primary key constraint error for table [{schema}.{table_name}]: {exc}"
+        )
+        return {"constrained_columns": []}
+
+
+def get_foreign_keys(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    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,
+            )
+        )
+        col_name = []
+        for table_constraint in table_constraints:
+            col_name.append(
+                {
+                    "name": table_constraint.name,
+                    "referred_schema": table_constraint.referred_schema,
+                    "referred_table": table_constraint.referred_table,
+                    "constrained_columns": [table_constraint.constrained_columns],
+                    "referred_columns": [table_constraint.referred_columns],
+                }
+            )
+        return col_name
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(
+            f"Error while fetching foreign key constraint error for table [{schema}.{table_name}]: {exc}"
+        )
+        return []