openmetadata-ingestion 1.7.5.0__py3-none-any.whl → 1.8.0.0__py3-none-any.whl

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

@@ -36,6 +36,12 @@ from metadata.ingestion.connections.test_connections import test_connection_step
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
 from metadata.ingestion.source.database.unitycatalog.models import DatabricksTable
+from metadata.ingestion.source.database.unitycatalog.queries import (
+    UNITY_CATALOG_GET_ALL_SCHEMA_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_TAGS,
+    UNITY_CATALOG_GET_CATALOGS_TAGS,
+)
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.db_utils import get_host_from_host_port
 from metadata.utils.logger import ingestion_logger
@@ -111,6 +117,32 @@ def test_connection(
                 table_obj.name = table.name
                 break
 
+    def get_tags(
+        service_connection: UnityCatalogConnection, table_obj: DatabricksTable
+    ):
+        engine = get_sqlalchemy_connection(service_connection)
+        with engine.connect() as connection:
+            connection.execute(
+                UNITY_CATALOG_GET_CATALOGS_TAGS.format(
+                    database=table_obj.catalog_name
+                ).replace(";", " limit 1;")
+            )
+            connection.execute(
+                UNITY_CATALOG_GET_ALL_SCHEMA_TAGS.format(
+                    database=table_obj.catalog_name
+                ).replace(";", " limit 1;")
+            )
+            connection.execute(
+                UNITY_CATALOG_GET_ALL_TABLE_TAGS.format(
+                    database=table_obj.catalog_name, schema=table_obj.schema_name
+                ).replace(";", " limit 1;")
+            )
+            connection.execute(
+                UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS.format(
+                    database=table_obj.catalog_name, schema=table_obj.schema_name
+                ).replace(";", " limit 1;")
+            )
+
     test_fn = {
         "CheckAccess": connection.catalogs.list,
         "GetDatabases": partial(get_catalogs, connection, table_obj),
@@ -118,6 +150,7 @@ def test_connection(
         "GetTables": partial(get_tables, connection, table_obj),
         "GetViews": partial(get_tables, connection, table_obj),
         "GetQueries": client.test_query_api_access,
+        "GetTags": partial(get_tags, service_connection, table_obj),
     }
 
     return test_connection_steps(

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

@@ -33,12 +33,9 @@ from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_test_connection_fn
+from metadata.ingestion.source.connections import test_connection_common
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
 from metadata.ingestion.source.database.unitycatalog.connection import get_connection
 from metadata.ingestion.source.database.unitycatalog.models import LineageTableStreams
@@ -220,8 +217,6 @@ class UnitycatalogLineageSource(Source):
                     )
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)

metadata/ingestion/source/database/unitycatalog/metadata.py

@@ -65,20 +65,33 @@ from metadata.ingestion.source.database.external_table_lineage_mixin import (
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
-from metadata.ingestion.source.database.unitycatalog.connection import get_connection
+from metadata.ingestion.source.database.unitycatalog.connection import (
+    get_connection,
+    get_sqlalchemy_connection,
+)
 from metadata.ingestion.source.database.unitycatalog.models import (
     ColumnJson,
     ElementType,
     ForeignConstrains,
     Type,
 )
+from metadata.ingestion.source.database.unitycatalog.queries import (
+    UNITY_CATALOG_GET_ALL_SCHEMA_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS,
+    UNITY_CATALOG_GET_ALL_TABLE_TAGS,
+    UNITY_CATALOG_GET_CATALOGS_TAGS,
+)
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
 from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
 logger = ingestion_logger()
 
+UNITY_CATALOG_TAG = "UNITY CATALOG TAG"
+UNITY_CATALOG_TAG_CLASSIFICATION = "UNITY CATALOG TAG CLASSIFICATION"
+
 
 class UnitycatalogSource(
     ExternalTableLineageMixin, DatabaseServiceSource, MultiDBSource
@@ -179,10 +192,14 @@ class UnitycatalogSource(
         From topology.
         Prepare a database request and pass it to the sink
         """
+        catalog = self.client.catalogs.get(database_name)
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
                 service=self.context.get().database_service,
+                owners=self.get_owner_ref(catalog.owner),
+                description=catalog.comment,
+                tags=self.get_database_tag_labels(database_name),
             )
         )
 
@@ -227,6 +244,9 @@ class UnitycatalogSource(
         From topology.
         Prepare a database schema request and pass it to the sink
         """
+        schema = self.client.schemas.get(
+            full_name=f"{self.context.get().database}.{schema_name}"
+        )
         yield Either(
             right=CreateDatabaseSchemaRequest(
                 name=EntityName(schema_name),
@@ -238,6 +258,9 @@ class UnitycatalogSource(
                         database_name=self.context.get().database,
                     )
                 ),
+                description=schema.comment,
+                owners=self.get_owner_ref(schema.owner),
+                tags=self.get_schema_tag_labels(schema_name),
             )
         )
 
@@ -315,7 +338,7 @@ class UnitycatalogSource(
                 (db_name, schema_name, table_name)
             ] = table.storage_location
         try:
-            columns = list(self.get_columns(table.columns))
+            columns = list(self.get_columns(table_name, table.columns))
             (
                 primary_constraints,
                 foreign_constraints,
@@ -341,6 +364,7 @@ class UnitycatalogSource(
                     )
                 ),
                 owners=self.get_owner_ref(table.owner),
+                tags=self.get_tag_labels(table_name),
             )
             yield Either(right=table_request)
 
@@ -477,11 +501,12 @@ class UnitycatalogSource(
                 f"Unable to add description to complex datatypes for column [{column.name}]: {exc}"
             )
 
-    def get_columns(self, column_data: List[ColumnInfo]) -> Iterable[Column]:
+    def get_columns(
+        self, table_name: str, column_data: List[ColumnInfo]
+    ) -> Iterable[Column]:
         """
         process table regular columns info
         """
-
         for column in column_data:
             parsed_string = {}
             if column.type_text:
@@ -500,6 +525,9 @@ class UnitycatalogSource(
             parsed_string["dataLength"] = parsed_string.get("dataLength", 1)
             if column.comment:
                 parsed_string["description"] = Markdown(column.comment)
+            parsed_string["tags"] = self.get_column_tag_labels(
+                table_name=table_name, column={"name": column.name}
+            )
             parsed_column = Column(**parsed_string)
             self.add_complex_datatype_descriptions(
                 column=parsed_column,
@@ -507,10 +535,97 @@ class UnitycatalogSource(
             )
             yield parsed_column
 
+    def yield_database_tag(
+        self, database_name: str
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """Get Unity Catalog database/catalog tags using SQL query"""
+        query_tag_fqn_builder_mapping = (
+            (
+                UNITY_CATALOG_GET_CATALOGS_TAGS.format(database=database_name),
+                lambda tag: [self.context.get().database_service, database_name],
+            ),
+            (
+                UNITY_CATALOG_GET_ALL_SCHEMA_TAGS.format(database=database_name),
+                lambda tag: [
+                    self.context.get().database_service,
+                    database_name,
+                    tag.schema_name,
+                ],
+            ),
+        )
+        try:
+            with get_sqlalchemy_connection(
+                self.service_connection
+            ).connect() as connection:
+                for query, tag_fqn_builder in query_tag_fqn_builder_mapping:
+                    for tag in connection.execute(query):
+                        yield from get_ometa_tag_and_classification(
+                            tag_fqn=FullyQualifiedEntityName(
+                                fqn._build(*tag_fqn_builder(tag))
+                            ),  # pylint: disable=protected-access
+                            tags=[tag.tag_value],
+                            classification_name=tag.tag_name,
+                            tag_description=UNITY_CATALOG_TAG,
+                            classification_description=UNITY_CATALOG_TAG_CLASSIFICATION,
+                            metadata=self.metadata,
+                            system_tags=True,
+                        )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error getting tags for catalog/schema {database_name}: {exc}"
+            )
+
     def yield_tag(
         self, schema_name: str
     ) -> Iterable[Either[OMetaTagAndClassification]]:
-        """No tags being processed"""
+        """Get Unity Catalog schema tags using SQL query"""
+        database = self.context.get().database
+        query_tag_fqn_builder_mapping = (
+            (
+                UNITY_CATALOG_GET_ALL_TABLE_TAGS.format(
+                    database=database, schema=schema_name
+                ),
+                lambda tag: [
+                    self.context.get().database_service,
+                    database,
+                    schema_name,
+                    tag.table_name,
+                ],
+            ),
+            (
+                UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS.format(
+                    database=database, schema=schema_name
+                ),
+                lambda tag: [
+                    self.context.get().database_service,
+                    database,
+                    schema_name,
+                    tag.table_name,
+                    tag.column_name,
+                ],
+            ),
+        )
+        try:
+            with get_sqlalchemy_connection(
+                self.service_connection
+            ).connect() as connection:
+                for query, tag_fqn_builder in query_tag_fqn_builder_mapping:
+                    for tag in connection.execute(query):
+                        yield from get_ometa_tag_and_classification(
+                            tag_fqn=FullyQualifiedEntityName(
+                                fqn._build(*tag_fqn_builder(tag))
+                            ),  # pylint: disable=protected-access
+                            tags=[tag.tag_value],
+                            classification_name=tag.tag_name,
+                            tag_description=UNITY_CATALOG_TAG,
+                            classification_description=UNITY_CATALOG_TAG_CLASSIFICATION,
+                            metadata=self.metadata,
+                            system_tags=True,
+                        )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error getting tags for schema {schema_name}: {exc}")
 
     def get_stored_procedures(self) -> Iterable[Any]:
         """Not implemented"""
@@ -527,24 +642,22 @@ class UnitycatalogSource(
         """Nothing to close"""
 
     # pylint: disable=arguments-renamed
-    def get_owner_ref(
-        self, table_owner: Optional[str]
-    ) -> Optional[EntityReferenceList]:
+    def get_owner_ref(self, owner: Optional[str]) -> Optional[EntityReferenceList]:
         """
         Method to process the table owners
         """
         if self.source_config.includeOwners is False:
             return None
         try:
-            if not table_owner or not isinstance(table_owner, str):
+            if not owner or not isinstance(owner, str):
                 return None
-            owner_ref = self.metadata.get_reference_by_email(email=table_owner)
+            owner_ref = self.metadata.get_reference_by_email(email=owner)
             if owner_ref:
                 return owner_ref
-            table_name = table_owner.split("@")[0]
-            owner_ref = self.metadata.get_reference_by_name(name=table_name)
+            owner_name = owner.split("@")[0]
+            owner_ref = self.metadata.get_reference_by_name(name=owner_name)
             return owner_ref
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Error processing owner {table_owner}: {exc}")
+            logger.warning(f"Error processing owner {owner}: {exc}")
         return None

metadata/ingestion/source/database/unitycatalog/queries.py

@@ -0,0 +1,19 @@
+"""
+SQL queries for Unity Catalog
+"""
+
+UNITY_CATALOG_GET_CATALOGS_TAGS = """
+SELECT * FROM `{database}`.information_schema.catalog_tags;
+"""
+
+UNITY_CATALOG_GET_ALL_SCHEMA_TAGS = """
+SELECT * FROM `{database}`.information_schema.schema_tags;
+"""
+
+UNITY_CATALOG_GET_ALL_TABLE_TAGS = """
+SELECT * FROM `{database}`.information_schema.table_tags WHERE schema_name = '{schema}';
+"""
+
+UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS = """
+SELECT * FROM `{database}`.information_schema.column_tags WHERE schema_name = '{schema}';
+"""

metadata/ingestion/source/messaging/messaging_service.py

@@ -34,9 +34,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
@@ -45,7 +42,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_topic
 from metadata.utils.helpers import retry_with_docker_host
@@ -202,11 +199,9 @@ class MessagingServiceSource(TopologyRunnerMixin, Source, ABC):
         """By default, nothing to prepare"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def mark_topics_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the topics as deleted"""

metadata/ingestion/source/metadata/alationsink/metadata.py

@@ -33,12 +33,9 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.models import Either, Entity
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.ometa.utils import model_str
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.metadata.alationsink.client import AlationSinkClient
 from metadata.ingestion.source.metadata.alationsink.constants import (
     SERVICE_TYPE_MAPPER,
@@ -475,8 +472,6 @@ class AlationsinkSource(Source):
         """Not required to implement"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.alation_sink_client, self.service_connection
         )
-        raise_test_connection_exception(result)

metadata/ingestion/source/metadata/amundsen/metadata.py

@@ -56,13 +56,10 @@ from metadata.generated.schema.type.entityReferenceList import EntityReferenceLi
 from metadata.ingestion.api.common import Entity
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.user import OMetaUserProfile
 from metadata.ingestion.ometa.client_utils import get_chart_entities_from_id
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.ingestion.source.metadata.amundsen.queries import (
     NEO4J_AMUNDSEN_DASHBOARD_QUERY,
@@ -465,8 +462,6 @@ class AmundsenSource(Source):
         return None
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)

metadata/ingestion/source/metadata/atlas/metadata.py

@@ -41,12 +41,9 @@ from metadata.generated.schema.type.entityLineage import EntitiesEdge
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either, Entity, StackTraceError
 from metadata.ingestion.api.steps import InvalidSourceException, Source
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.ingestion.source.metadata.atlas.client import AtlasClient
 from metadata.utils import fqn
@@ -480,8 +477,6 @@ class AtlasSource(Source):
         return None
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -38,9 +38,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
@@ -49,7 +46,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.utils import fqn
 from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
@@ -180,11 +177,9 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
         """By default, nothing to close"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def mark_mlmodels_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the mlmodels as deleted"""

metadata/ingestion/source/pipeline/airflow/connection.py

@@ -24,7 +24,7 @@ from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
 from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
-    MysqlConnection,
+    MysqlConnection as MysqlConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
     PostgresConnection,
@@ -67,12 +67,10 @@ def _(_: BackendConnection) -> Engine:
 
 
 @_get_connection.register
-def _(airflow_connection: MysqlConnection) -> Engine:
-    from metadata.ingestion.source.database.mysql.connection import (
-        get_connection as get_mysql_connection,
-    )
+def _(airflow_connection: MysqlConnectionConfig) -> Engine:
+    from metadata.ingestion.source.database.mysql.connection import MySQLConnection
 
-    return get_mysql_connection(airflow_connection)
+    return MySQLConnection(airflow_connection).get_client()
 
 
 @_get_connection.register