openmetadata-ingestion 1.2.5.2__py3-none-any.whl → 1.3.0.0__py3-none-any.whl

metadata/ingestion/source/database/snowflake/metadata.py

@@ -13,7 +13,6 @@ Snowflake source module
 """
 import json
 import traceback
-import urllib
 from typing import Dict, Iterable, List, Optional, Tuple
 
 import sqlparse
@@ -30,21 +29,21 @@ from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.data.table import (
     IntervalType,
-    Table,
     TablePartition,
     TableType,
 )
 from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
     SnowflakeConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.type.basic import EntityName, SourceUrl
-from metadata.generated.schema.type.lifeCycle import AccessDetails, LifeCycle
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
-from metadata.ingestion.models.life_cycle import OMetaLifeCycleData
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.column_type_parser import create_sqlalchemy_type
@@ -99,7 +98,6 @@ from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import get_all_table_comments
 from metadata.utils.tag_utils import get_ometa_tag_and_classification
-from metadata.utils.time_utils import convert_timestamp_to_milliseconds
 
 ischema_names["VARIANT"] = VARIANT
 ischema_names["GEOGRAPHY"] = create_sqlalchemy_type("GEOGRAPHY")
@@ -145,6 +143,7 @@ class SnowflakeSource(
 
         self._account: Optional[str] = None
         self._org_name: Optional[str] = None
+        self.life_cycle_query = SNOWFLAKE_LIFE_CYCLE_QUERY
 
     @classmethod
     def create(cls, config_dict, metadata: OpenMetadata):
@@ -363,7 +362,7 @@ class SnowflakeSource(
                         left=StackTraceError(
                             name="Tags and Classifications",
                             error=f"Failed to fetch tags due to [{inner_exc}]",
-                            stack_trace=traceback.format_exc(),
+                            stackTrace=traceback.format_exc(),
                         )
                     )
 
@@ -476,48 +475,6 @@ class SnowflakeSource(
             logger.error(f"Unable to get source url: {exc}")
         return None
 
-    def yield_life_cycle_data(self, _) -> Iterable[Either[OMetaLifeCycleData]]:
-        """
-        Get the life cycle data of the table
-        """
-        table_fqn = fqn.build(
-            self.metadata,
-            entity_type=Table,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
-            schema_name=self.context.database_schema,
-            table_name=self.context.table,
-            skip_es_search=True,
-        )
-        table = self.metadata.get_by_name(entity=Table, fqn=table_fqn)
-        if table:
-            try:
-                life_cycle_data = self.life_cycle_query_dict(
-                    query=SNOWFLAKE_LIFE_CYCLE_QUERY.format(
-                        database_name=table.database.name,
-                        schema_name=table.databaseSchema.name,
-                    )
-                ).get(table.name.__root__)
-                if life_cycle_data:
-                    life_cycle = LifeCycle(
-                        created=AccessDetails(
-                            timestamp=convert_timestamp_to_milliseconds(
-                                life_cycle_data.created_at.timestamp()
-                            )
-                        )
-                    )
-                    yield Either(
-                        right=OMetaLifeCycleData(entity=table, life_cycle=life_cycle)
-                    )
-            except Exception as exc:
-                yield Either(
-                    left=StackTraceError(
-                        name=table.name.__root__,
-                        error=f"Unable to get the table life cycle data for table {table.name.__root__}: {exc}",
-                        stack_trace=traceback.format_exc(),
-                    )
-                )
-
     def query_view_names_and_types(
         self, schema_name: str
     ) -> Iterable[TableNameAndType]:
@@ -581,7 +538,7 @@ class SnowflakeSource(
                 database_name=self.context.database,
                 schema_name=self.context.database_schema,
                 procedure_name=stored_procedure.name,
-                procedure_signature=urllib.parse.unquote(stored_procedure.signature),
+                procedure_signature=stored_procedure.unquote_signature(),
             )
         )
         return dict(res.all()).get("body", "")
@@ -623,7 +580,7 @@ class SnowflakeSource(
                 left=StackTraceError(
                     name=stored_procedure.name,
                     error=f"Error yielding Stored Procedure [{stored_procedure.name}] due to [{exc}]",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

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

@@ -11,6 +11,7 @@
 """
 Snowflake models
 """
+import urllib
 from typing import Optional
 
 from pydantic import BaseModel, Field, validator
@@ -58,7 +59,8 @@ class SnowflakeStoredProcedure(BaseModel):
         try:
             clean_signature = signature.replace("(", "").replace(")", "")
             if not clean_signature:
-                return None
+                # If removing the () leaves us with nothing, then just return the parenthesis
+                return "()"
 
             signature_list = clean_signature.split(",")
             clean_signature_list = [elem.split(" ")[-1] for elem in signature_list]
@@ -67,3 +69,6 @@ class SnowflakeStoredProcedure(BaseModel):
         except Exception as exc:
             logger.warning(f"Error cleaning up Stored Procedure signature - [{exc}]")
             return signature
+
+    def unquote_signature(self) -> Optional[str]:
+        return urllib.parse.unquote(self.signature) if self.signature else "()"

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

@@ -187,7 +187,6 @@ SNOWFLAKE_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent(
     """
 WITH SP_HISTORY AS (
     SELECT
-      QUERY_ID,
       QUERY_TEXT,
       SESSION_ID,
       START_TIME,
@@ -198,7 +197,6 @@ WITH SP_HISTORY AS (
 ),
 Q_HISTORY AS (
     SELECT
-      QUERY_ID,
       QUERY_TYPE,
       QUERY_TEXT,
       SESSION_ID,
@@ -215,8 +213,6 @@ Q_HISTORY AS (
       AND START_TIME >= '{start_date}' 
 )
 SELECT
-  SP.QUERY_ID AS PROCEDURE_ID,
-  Q.QUERY_ID AS QUERY_ID,
   Q.QUERY_TYPE AS QUERY_TYPE,
   Q.DATABASE_NAME AS QUERY_DATABASE_NAME,
   Q.SCHEMA_NAME AS QUERY_SCHEMA_NAME,

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

@@ -24,7 +24,6 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
 from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_SESSION_TAG_QUERY,
@@ -61,7 +60,9 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
             filters=self.get_filters(),
         )
 
-    def check_life_cycle_query(self, query_type: Optional[str]) -> bool:
+    def check_life_cycle_query(
+        self, query_type: Optional[str], query_text: Optional[str]
+    ) -> bool:
         """
         returns true if query is to be used for life cycle processing.
 
@@ -87,24 +88,8 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
             )
 
     def get_table_query(self) -> Iterable[TableQuery]:
-        database = self.config.serviceConnection.__root__.config.database
-        if database:
-            use_db_query = f"USE DATABASE {database}"
-            self.engine.execute(use_db_query)
-            self.set_session_query_tag()
-            yield from super().get_table_query()
-        else:
-            query = "SHOW DATABASES"
-            results = self.engine.execute(query)
-            for res in results:
-                row = list(res)
-                use_db_query = f"USE DATABASE {row[1]}"
-                self.engine.execute(use_db_query)
-                logger.info(f"Ingesting from database: {row[1]}")
-                self.config.serviceConnection.__root__.config.database = row[1]
-                self.engine = get_connection(self.service_connection)
-                self.set_session_query_tag()
-                yield from super().get_table_query()
+        self.set_session_query_tag()
+        yield from super().get_table_query()
 
     def get_database_name(self, data: dict) -> str:  # pylint: disable=arguments-differ
         """

metadata/ingestion/source/database/snowflake/utils.py

@@ -39,7 +39,6 @@ from metadata.utils.sqlalchemy_utils import (
 
 
 def _quoted_name(entity_name: Optional[str]) -> Optional[str]:
-
     if entity_name:
         return fqn.quote_name(entity_name)
 
@@ -121,13 +120,10 @@ def get_view_definition(  # pylint: disable=unused-argument
     schema = schema or self.default_schema_name
     if schema:
         cursor = connection.execute(
-            "SHOW /* sqlalchemy:get_view_definition */ VIEWS "
-            f"LIKE '{view_name}' IN {schema}"
+            f"SELECT GET_DDL('VIEW','{schema}.{view_name}') AS \"text\""
         )
     else:
-        cursor = connection.execute(
-            "SHOW /* sqlalchemy:get_view_definition */ VIEWS " f"LIKE '{view_name}'"
-        )
+        cursor = connection.execute(f"SELECT GET_DDL('VIEW','{view_name}') AS \"text\"")
     n2i = self.__class__._map_name_to_idx(cursor)  # pylint: disable=protected-access
     try:
         ret = cursor.fetchone()

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -24,13 +24,16 @@ from sqlalchemy.engine import Engine
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
 from metadata.generated.schema.type.basic import SqlQuery, Timestamp
 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, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.status import Status
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
@@ -48,15 +51,14 @@ class QueryByProcedure(BaseModel):
     Query(ies) executed by each stored procedure
     """
 
-    procedure_id: str = Field(..., alias="PROCEDURE_ID")
-    query_id: str = Field(..., alias="QUERY_ID")
+    procedure_name: str = Field(None, alias="PROCEDURE_NAME")
     query_type: str = Field(..., alias="QUERY_TYPE")
     query_database_name: str = Field(None, alias="QUERY_DATABASE_NAME")
     query_schema_name: str = Field(None, alias="QUERY_SCHEMA_NAME")
     procedure_text: str = Field(..., alias="PROCEDURE_TEXT")
     procedure_start_time: datetime = Field(..., alias="PROCEDURE_START_TIME")
     procedure_end_time: datetime = Field(..., alias="PROCEDURE_END_TIME")
-    query_start_time: datetime = Field(..., alias="QUERY_START_TIME")
+    query_start_time: Optional[datetime] = Field(..., alias="QUERY_START_TIME")
     query_duration: Optional[float] = Field(None, alias="QUERY_DURATION")
     query_text: str = Field(..., alias="QUERY_TEXT")
     query_user_name: Optional[str] = Field(None, alias="QUERY_USER_NAME")
@@ -106,8 +108,11 @@ class StoredProcedureMixin(ABC):
         for row in results:
             try:
                 query_by_procedure = QueryByProcedure.parse_obj(dict(row))
-                procedure_name = get_procedure_name_from_call(
-                    query_text=query_by_procedure.procedure_text,
+                procedure_name = (
+                    query_by_procedure.procedure_name
+                    or get_procedure_name_from_call(
+                        query_text=query_by_procedure.procedure_text,
+                    )
                 )
                 queries_dict[procedure_name].append(query_by_procedure)
             except Exception as exc:
@@ -115,7 +120,7 @@ class StoredProcedureMixin(ABC):
                     StackTraceError(
                         name="Stored Procedure",
                         error=f"Error trying to get procedure name due to [{exc}]",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -148,7 +153,6 @@ class StoredProcedureMixin(ABC):
             query_type=query_by_procedure.query_type,
             query_text=query_by_procedure.query_text,
         ):
-
             self.context.stored_procedure_query_lineage = True
             for either_lineage in get_lineage_by_query(
                 self.metadata,

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

@@ -0,0 +1,87 @@
+#  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.
+"""
+Client to interact with databricks apis
+"""
+import json
+import traceback
+
+from metadata.ingestion.source.database.databricks.client import (
+    API_TIMEOUT,
+    DatabricksClient,
+)
+from metadata.ingestion.source.database.unitycatalog.models import (
+    LineageColumnStreams,
+    LineageTableStreams,
+)
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+TABLE_LINEAGE_PATH = "/lineage-tracking/table-lineage/get"
+COLUMN_LINEAGE_PATH = "/lineage-tracking/column-lineage/get"
+
+
+class UnityCatalogClient(DatabricksClient):
+    """
+    UnityCatalogClient creates a Databricks connection based on DatabricksCredentials.
+    """
+
+    def get_table_lineage(self, table_name: str) -> LineageTableStreams:
+        """
+        Method returns table lineage details
+        """
+        try:
+            data = {
+                "table_name": table_name,
+            }
+
+            response = self.client.get(
+                f"{self.base_url}{TABLE_LINEAGE_PATH}",
+                headers=self.headers,
+                data=json.dumps(data),
+                timeout=API_TIMEOUT,
+            ).json()
+            if response:
+                return LineageTableStreams(**response)
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(exc)
+
+        return LineageTableStreams()
+
+    def get_column_lineage(
+        self, table_name: str, column_name: str
+    ) -> LineageColumnStreams:
+        """
+        Method returns table lineage details
+        """
+        try:
+            data = {
+                "table_name": table_name,
+                "column_name": column_name,
+            }
+
+            response = self.client.get(
+                f"{self.base_url}{COLUMN_LINEAGE_PATH}",
+                headers=self.headers,
+                data=json.dumps(data),
+                timeout=API_TIMEOUT,
+            ).json()
+
+            if response:
+                return LineageColumnStreams(**response)
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(exc)
+
+        return LineageColumnStreams()

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

@@ -0,0 +1,97 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Source connection handler
+"""
+from functools import partial
+from typing import Optional
+
+from databricks.sdk import WorkspaceClient
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
+from metadata.ingestion.connections.test_connections import test_connection_steps
+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.utils.db_utils import get_host_from_host_port
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+def get_connection_url(connection: UnityCatalogConnection) -> str:
+    url = f"{connection.scheme.value}://token:{connection.token.get_secret_value()}@{connection.hostPort}"
+    return url
+
+
+def get_connection(connection: UnityCatalogConnection) -> WorkspaceClient:
+    """
+    Create connection
+    """
+
+    return WorkspaceClient(
+        host=get_host_from_host_port(connection.hostPort),
+        token=connection.token.get_secret_value(),
+    )
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    connection: WorkspaceClient,
+    service_connection: UnityCatalogConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+) -> None:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+    client = UnityCatalogClient(service_connection)
+    table_obj = DatabricksTable()
+
+    def get_catalogs(connection: WorkspaceClient, table_obj: DatabricksTable):
+        for catalog in connection.catalogs.list():
+            table_obj.catalog_name = catalog.name
+            break
+
+    def get_schemas(connection: WorkspaceClient, table_obj: DatabricksTable):
+        for schema in connection.schemas.list(catalog_name=table_obj.catalog_name):
+            table_obj.schema_name = schema.name
+            break
+
+    def get_tables(connection: WorkspaceClient, table_obj: DatabricksTable):
+        for table in connection.tables.list(
+            catalog_name=table_obj.catalog_name, schema_name=table_obj.schema_name
+        ):
+            table_obj.name = table.name
+            break
+
+    test_fn = {
+        "CheckAccess": connection.catalogs.list,
+        "GetDatabases": partial(get_catalogs, connection, table_obj),
+        "GetSchemas": partial(get_schemas, connection, table_obj),
+        "GetTables": partial(get_tables, connection, table_obj),
+        "GetViews": partial(get_tables, connection, table_obj),
+        "GetQueries": client.test_query_api_access,
+    }
+
+    test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+        timeout_seconds=service_connection.connectionTimeout,
+    )

metadata/ingestion/source/database/{databricks/unity_catalog → unitycatalog}/lineage.py

@@ -16,8 +16,8 @@ from typing import Iterable, Optional
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.table import Table
-from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
-    DatabricksConnection,
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
@@ -33,18 +33,18 @@ from metadata.ingestion.api.steps import InvalidSourceException, Source
 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.database.databricks.client import DatabricksClient
-from metadata.ingestion.source.database.databricks.connection import get_connection
-from metadata.ingestion.source.database.databricks.models import LineageTableStreams
+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
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
-class DatabricksUnityCatalogLineageSource(Source):
+class UnitycatalogLineageSource(Source):
     """
-    Databricks Lineage Unity Catalog Source
+    Lineage Unity Catalog Source
     """
 
     def __init__(
@@ -57,7 +57,7 @@ class DatabricksUnityCatalogLineageSource(Source):
         self.metadata = metadata
         self.service_connection = self.config.serviceConnection.__root__.config
         self.source_config = self.config.sourceConfig.config
-        self.client = DatabricksClient(self.service_connection)
+        self.client = UnityCatalogClient(self.service_connection)
         self.connection_obj = get_connection(self.service_connection)
         self.test_connection()
 
@@ -75,10 +75,10 @@ class DatabricksUnityCatalogLineageSource(Source):
     def create(cls, config_dict, metadata: OpenMetadata):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: DatabricksConnection = config.serviceConnection.__root__.config
-        if not isinstance(connection, DatabricksConnection):
+        connection: UnityCatalogConnection = config.serviceConnection.__root__.config
+        if not isinstance(connection, UnityCatalogConnection):
             raise InvalidSourceException(
-                f"Expected DatabricksConnection, but got {connection}"
+                f"Expected UnityCatalogConnection, but got {connection}"
             )
         return cls(config, metadata)