openmetadata-ingestion 1.4.0.0rc1__py3-none-any.whl → 1.4.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/mysql/metadata.py

@@ -13,6 +13,7 @@ from typing import Optional, cast
 
 from sqlalchemy.dialects.mysql.base import ischema_names
 from sqlalchemy.dialects.mysql.reflection import MySQLTableDefinitionParser
+from sqlalchemy.engine.reflection import Inspector
 
 from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
     MysqlConnection,
@@ -24,6 +25,7 @@ from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.common_db_source import CommonDbSourceService
 from metadata.ingestion.source.database.mysql.utils import col_type_map, parse_column
+from metadata.utils.sqlalchemy_utils import get_all_table_ddls, get_table_ddl
 
 ischema_names.update(col_type_map)
 
@@ -32,6 +34,9 @@ MySQLTableDefinitionParser._parse_column = (  # pylint: disable=protected-access
     parse_column
 )
 
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
+
 
 class MysqlSource(CommonDbSourceService):
     """

metadata/ingestion/source/database/oracle/metadata.py

@@ -71,7 +71,9 @@ 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,
+    get_all_table_ddls,
     get_all_view_definitions,
+    get_table_ddl,
 )
 
 logger = ingestion_logger()
@@ -96,6 +98,9 @@ Inspector.get_mview_names = get_mview_names
 Inspector.get_mview_definition = get_mview_definition
 OracleDialect.get_mview_names = get_mview_names_dialect
 
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
+
 
 class OracleSource(StoredProcedureMixin, CommonDbSourceService):
     """
@@ -138,27 +143,36 @@ class OracleSource(StoredProcedureMixin, CommonDbSourceService):
 
         return regular_tables + material_tables
 
-    def get_view_definition(
+    def get_schema_definition(
         self, table_type: str, table_name: str, schema_name: str, inspector: Inspector
     ) -> Optional[str]:
-        if table_type not in {TableType.View, TableType.MaterializedView}:
-            return None
+        """
+        Get the DDL statement or View Definition for a table
+        """
+        try:
+            if table_type not in {TableType.View, TableType.MaterializedView}:
+                schema_definition = inspector.get_table_ddl(
+                    self.connection, table_name, schema_name
+                )
 
-        definition_fn = inspector.get_view_definition
-        if table_type == TableType.MaterializedView:
-            definition_fn = inspector.get_mview_definition
+            else:
+                definition_fn = inspector.get_view_definition
+                if table_type == TableType.MaterializedView:
+                    definition_fn = inspector.get_mview_definition
 
-        try:
-            view_definition = definition_fn(table_name, schema_name)
-            view_definition = "" if view_definition is None else str(view_definition)
-            return view_definition
+                schema_definition = definition_fn(table_name, schema_name)
+
+            schema_definition = (
+                str(schema_definition) if schema_definition is not None else None
+            )
+            return schema_definition
 
         except NotImplementedError:
-            logger.warning("View definition not implemented")
+            logger.warning("Schema definition not implemented")
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Failed to fetch view definition for {table_name}: {exc}")
+            logger.warning(f"Failed to fetch Schema definition for {table_name}: {exc}")
         return None
 
     def process_result(self, data: FetchProcedureList):

metadata/ingestion/source/database/postgres/metadata.py

@@ -66,8 +66,10 @@ from metadata.utils.filters import filter_by_database
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
     get_all_table_comments,
+    get_all_table_ddls,
     get_all_table_owners,
     get_all_view_definitions,
+    get_table_ddl,
 )
 from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
@@ -124,6 +126,8 @@ PGDialect.get_all_table_owners = get_all_table_owners
 PGDialect.get_table_owner = get_table_owner
 PGDialect.ischema_names = ischema_names
 
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
 Inspector.get_table_owner = get_etable_owner
 
 PGDialect.get_foreign_keys = get_foreign_keys

metadata/ingestion/source/database/redshift/metadata.py

@@ -96,7 +96,11 @@ from metadata.utils.execution_time_tracker import (
 from metadata.utils.filters import filter_by_database
 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.sqlalchemy_utils import (
+    get_all_table_comments,
+    get_all_table_ddls,
+    get_table_ddl,
+)
 
 logger = ingestion_logger()
 
@@ -122,6 +126,9 @@ RedshiftDialect._get_all_relation_info = (  # pylint: disable=protected-access
     _get_all_relation_info
 )
 
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
+
 
 class RedshiftSource(
     LifeCycleQueryMixin, StoredProcedureMixin, CommonDbSourceService, MultiDBSource

metadata/ingestion/source/database/sample_data.py

@@ -733,6 +733,7 @@ class SampleDataSource(
                 tableType=table["tableType"],
                 tableConstraints=table.get("tableConstraints"),
                 tags=table["tags"],
+                schemaDefinition=table.get("schemaDefinition"),
             )
 
             yield Either(right=table_and_db)
@@ -1567,9 +1568,6 @@ class SampleDataSource(
     def ingest_life_cycle(self) -> Iterable[Either[OMetaLifeCycleData]]:
         """Iterate over all the life cycle data and ingest them"""
         for table_life_cycle in self.life_cycle_data["lifeCycleData"]:
-            table = self.metadata.get_by_name(
-                entity=Table, fqn=table_life_cycle["fqn"], fields=["lifeCycle"]
-            )
             life_cycle = table_life_cycle["lifeCycle"]
             life_cycle_data = LifeCycle()
             life_cycle_data.created = AccessDetails(

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

@@ -92,6 +92,7 @@ from metadata.ingestion.source.database.snowflake.utils import (
     get_schema_columns,
     get_schema_foreign_keys,
     get_table_comment,
+    get_table_ddl,
     get_table_names,
     get_table_names_reflection,
     get_unique_constraints,
@@ -108,7 +109,7 @@ from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database
 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.sqlalchemy_utils import get_all_table_comments, get_all_table_ddls
 from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
 ischema_names["VARIANT"] = VARIANT
@@ -137,6 +138,8 @@ SnowflakeDialect._current_database_schema = (  # pylint: disable=protected-acces
 SnowflakeDialect.get_pk_constraint = get_pk_constraint
 SnowflakeDialect.get_foreign_keys = get_foreign_keys
 SnowflakeDialect.get_columns = get_columns
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
 SnowflakeDialect._get_schema_foreign_keys = get_schema_foreign_keys
 
 

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

@@ -351,3 +351,7 @@ JOIN Q_HISTORY Q
 ORDER BY PROCEDURE_START_TIME DESC
     """
 )
+
+SNOWFLAKE_GET_TABLE_DDL = """
+SELECT GET_DDL('TABLE','{table_name}') AS \"text\"
+"""

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

@@ -35,6 +35,7 @@ from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_GET_EXTERNAL_TABLE_NAMES,
     SNOWFLAKE_GET_MVIEW_NAMES,
     SNOWFLAKE_GET_SCHEMA_COLUMNS,
+    SNOWFLAKE_GET_TABLE_DDL,
     SNOWFLAKE_GET_TRANSIENT_NAMES,
     SNOWFLAKE_GET_VIEW_NAMES,
     SNOWFLAKE_GET_WITHOUT_TRANSIENT_TABLE_NAMES,
@@ -476,3 +477,22 @@ def get_columns(self, connection, table_name, schema=None, **kw):
     if normalized_table_name not in schema_columns:
         raise sa_exc.NoSuchTableError()
     return schema_columns[normalized_table_name]
+
+
+@reflection.cache
+def get_table_ddl(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    Gets the Table DDL
+    """
+    schema = schema or self.default_schema_name
+    table_name = f"{schema}.{table_name}" if schema else table_name
+    cursor = connection.execute(SNOWFLAKE_GET_TABLE_DDL.format(table_name=table_name))
+    try:
+        result = cursor.fetchone()
+        if result:
+            return result[0]
+    except Exception:
+        pass
+    return None

metadata/ingestion/source/database/sqlalchemy_source.py

@@ -74,11 +74,11 @@ class SqlAlchemySource(ABC):
         """
 
     @abstractmethod
-    def get_view_definition(
+    def get_schema_definition(
         self, table_type, table_name: str, schema_name: str, inspector: Inspector
     ) -> Optional[str]:
         """
-        Method to fetch view definition
+        Method to fetch schema definition
         """
 
     @abstractmethod

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

@@ -18,6 +18,7 @@ from urllib.parse import quote_plus
 from requests import Session
 from sqlalchemy.engine import Engine
 
+from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
@@ -98,6 +99,18 @@ def get_connection(connection: TrinoConnection) -> Engine:
         connection.connectionArguments.__root__["verify"] = {
             "verify": connection.verify
         }
+    if hasattr(connection.authType, "azureConfig"):
+        azure_client = AzureClient(connection.authType.azureConfig).create_client()
+        if not connection.authType.azureConfig.scopes:
+            raise ValueError(
+                "Azure Scopes are missing, please refer https://learn.microsoft.com/en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-microsoft-entra-access-token and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
+            )
+        access_token_obj = azure_client.get_token(
+            *connection.authType.azureConfig.scopes.split(",")
+        )
+        if not connection.connectionOptions:
+            connection.connectionOptions = init_empty_connection_options()
+        connection.connectionOptions.__root__["access_token"] = access_token_obj.token
     return create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url,

metadata/ingestion/source/messaging/common_broker_source.py

@@ -47,6 +47,7 @@ from metadata.parsers.schema_parsers import (
 )
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.messaging_utils import merge_and_clean_protobuf_schema
 
 logger = ingestion_logger()
 
@@ -75,6 +76,7 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
         self.generate_sample_data = self.config.sourceConfig.config.generateSampleData
         self.admin_client = self.connection.admin_client
         self.schema_registry_client = self.connection.schema_registry_client
+        self.context.processed_schemas = {}
         if self.generate_sample_data:
             self.consumer_client = self.connection.consumer_client
 
@@ -125,9 +127,14 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
                     raise InvalidSchemaTypeException(
                         f"Cannot find {schema_type} in parser providers registry."
                     )
-                schema_fields = load_parser_fn(
-                    topic_details.topic_name, topic_schema.schema_str
-                )
+                schema_text = topic_schema.schema_str
+
+                # In protobuf schema, we need to merge all the schema text with references
+                if schema_type == SchemaType.Protobuf.value.lower():
+                    schema_text = merge_and_clean_protobuf_schema(
+                        self._get_schema_text_with_references(schema=topic_schema)
+                    )
+                schema_fields = load_parser_fn(topic_details.topic_name, schema_text)
 
                 topic.messageSchema = Topic(
                     schemaText=topic_schema.schema_str,
@@ -195,6 +202,38 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
                 f"Exception adding properties to topic [{topic.name}]: {exc}"
             )
 
+    def _get_schema_text_with_references(self, schema) -> Optional[str]:
+        """
+        Returns the schema text with references resolved using recursive calls
+        """
+        try:
+            if schema:
+                schema_text = schema.schema_str
+                for reference in schema.references or []:
+                    if not self.context.processed_schemas.get(reference.name):
+                        self.context.processed_schemas[reference.name] = True
+                        reference_schema = (
+                            self.schema_registry_client.get_latest_version(
+                                reference.name
+                            )
+                        )
+                        if reference_schema.schema.references:
+                            schema_text = (
+                                schema_text
+                                + self._get_schema_text_with_references(
+                                    reference_schema.schema
+                                )
+                            )
+                        else:
+                            schema_text = (
+                                schema_text + reference_schema.schema.schema_str
+                            )
+                return schema_text
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to get schema with references: {exc}")
+        return None
+
     def _parse_topic_metadata(self, topic_name: str) -> Optional[Schema]:
         try:
             if self.schema_registry_client:

metadata/ingestion/source/pipeline/dagster/queries.py

@@ -14,62 +14,30 @@ GraphQL queries used during ingestion
 """
 
 DAGSTER_PIPELINE_DETAILS_GRAPHQL = """
-            query PipelineRuns{
-  repositoriesOrError{
-    __typename 
-    ... on RepositoryConnection{
-      nodes{
-        id,
-        name,
+query PipelineRuns {
+  repositoriesOrError {
+    __typename
+    ... on RepositoryConnection {
+      nodes {
+        id
+        name
         location {
           id
           name
         }
-        pipelines{
-          id,
-          name,
-          description,
-          runs{
-            id,
-            runId,
-            status,
-            startTime,
-            stats{
-                ... on RunStatsSnapshot {
-                        startTime
-                        endTime
-                        stepsFailed
-                    }
-                }
-            executionPlan{
-              steps{
-                inputs{
-                  name,
-                  dependsOn{
-                    inputs{
-                      name
-                    },
-                    outputs{
-                      name
-                    }
-                  }
-                },
-                outputs{
-                  name
-                }
-              }
-
-            }
-          }
+        pipelines {
+          id
+          name
+          description
         }
       }
     }
-
   }
 }
 """
 
-GRAPHQL_RUNS_QUERY = """query SidebarOpGraphsQuery($selector: PipelineSelector!, $handleID: String!) {
+GRAPHQL_RUNS_QUERY = """
+query SidebarOpGraphsQuery($selector: PipelineSelector!, $handleID: String!) {
   pipelineOrError(params: $selector) {
     __typename
     ... on Pipeline {
@@ -100,41 +68,44 @@ GRAPHQL_RUNS_QUERY = """query SidebarOpGraphsQuery($selector: PipelineSelector!,
 }
 """
 
-TEST_QUERY_GRAPHQL = """query Pipeline {
-            pipelineRunsOrError {
-                __typename
-                ... on PipelineRuns {
-                results {
-                    assetSelection {
-                    path
-                    }
-                    runId
-                    status
-                }
-                }
-            }
-            }"""
-
-GRAPHQL_QUERY_FOR_JOBS = """query PipelineRuns($selector: GraphSelector!){
-                graphOrError(selector: $selector){
-                    ... on Graph{
-                    id,
-                    name,
-                    description,
-                        solidHandles{
-                            handleID,
-                            solid{
-                                name,
-                                inputs{
-                                    dependsOn {
-                                        solid {
-                                            name
-                                        }
-                                    }
-                                },
-                            }
-                        },
+TEST_QUERY_GRAPHQL = """
+query Pipeline {
+  pipelineRunsOrError(limit: 1) {
+    __typename
+    ... on PipelineRuns {
+      results {
+        assetSelection {
+          path
+        }
+        runId
+        status
+      }
+    }
+  }
+}
+"""
 
-                    }
-                }
-                }"""
+GRAPHQL_QUERY_FOR_JOBS = """
+query PipelineRuns($selector: GraphSelector!) {
+  graphOrError(selector: $selector) {
+    ... on Graph {
+      id
+      name
+      description
+      solidHandles {
+        handleID
+        solid {
+          name
+          inputs {
+            dependsOn {
+              solid {
+                name
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+"""