openmetadata-ingestion 1.9.13.0__py3-none-any.whl → 1.10.1.0__py3-none-any.whl

metadata/ingestion/source/database/mariadb/metadata.py

@@ -11,21 +11,44 @@
 """
 MariaDB source module
 """
-from typing import Optional
+import traceback
+from typing import Iterable, Optional
 
 from sqlalchemy.dialects.mysql.base import ischema_names
 from sqlalchemy.dialects.mysql.reflection import MySQLTableDefinitionParser
 
+from metadata.generated.schema.api.data.createStoredProcedure import (
+    CreateStoredProcedureRequest,
+)
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
     MariaDBConnection,
 )
+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, Markdown
+from metadata.ingestion.api.models import Either
 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.mariadb.models import (
+    ROUTINE_TYPE_MAP,
+    MariaDBStoredProcedure,
+)
+from metadata.ingestion.source.database.mariadb.queries import (
+    MARIADB_GET_FUNCTIONS,
+    MARIADB_GET_STORED_PROCEDURES,
+)
 from metadata.ingestion.source.database.mysql.utils import col_type_map, parse_column
+from metadata.utils import fqn
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 ischema_names.update(col_type_map)
 
@@ -52,3 +75,65 @@ class MariadbSource(CommonDbSourceService):
                 f"Expected MariaDBConnection, but got {connection}"
             )
         return cls(config, metadata)
+
+    def _get_stored_procedures_internal(
+        self, query: str
+    ) -> Iterable[MariaDBStoredProcedure]:
+        results = self.engine.execute(query).all()
+        for row in results:
+            try:
+                yield MariaDBStoredProcedure.model_validate(dict(row._mapping))
+            except Exception as exc:
+                error = f"Error parsing Stored Procedure payload: {exc}"
+                logger.error(error)
+                self.status.failed(
+                    error=StackTraceError(
+                        name=dict(row).get("procedure_name", "UNKNOWN"),
+                        error=error,
+                        stackTrace=traceback.format_exc(),
+                    )
+                )
+
+    def get_stored_procedures(self) -> Iterable[MariaDBStoredProcedure]:
+        """List stored procedures"""
+        if self.source_config.includeStoredProcedures:
+            for query in (MARIADB_GET_STORED_PROCEDURES, MARIADB_GET_FUNCTIONS):
+                yield from self._get_stored_procedures_internal(
+                    query.format(schema_name=self.context.get().database_schema)
+                )
+
+    def yield_stored_procedure(
+        self, stored_procedure
+    ) -> Iterable[Either[CreateStoredProcedureRequest]]:
+        """Prepare the stored procedure payload"""
+        try:
+            stored_procedure_request = CreateStoredProcedureRequest(
+                name=EntityName(stored_procedure.name),
+                description=(
+                    Markdown(stored_procedure.description)
+                    if stored_procedure.description
+                    else None
+                ),
+                storedProcedureCode=StoredProcedureCode(
+                    language=stored_procedure.language, code=stored_procedure.definition
+                ),
+                databaseSchema=fqn.build(
+                    metadata=self.metadata,
+                    entity_type=DatabaseSchema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
+                ),
+                storedProcedureType=ROUTINE_TYPE_MAP[stored_procedure.procedure_type],
+            )
+            yield Either(right=stored_procedure_request)
+            self.register_record_stored_proc_request(stored_procedure_request)
+
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name=stored_procedure.name,
+                    error=f"Error yielding Stored Procedure [{stored_procedure.name}] due to [{exc}]",
+                    stackTrace=traceback.format_exc(),
+                )
+            )

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

@@ -0,0 +1,36 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  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.
+"""
+MariaDB models
+"""
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureType
+
+ROUTINE_TYPE_MAP = {
+    "PROCEDURE": StoredProcedureType.StoredProcedure,
+    "FUNCTION": StoredProcedureType.Function,
+}
+
+
+class MariaDBStoredProcedure(BaseModel):
+    """
+    MariaDB stored procedure list query results
+    """
+
+    name: str = Field(alias="procedure_name")
+    schema_name: str
+    definition: str
+    language: Optional[str]
+    procedure_type: Optional[str]
+    description: Optional[str]

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

@@ -0,0 +1,43 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  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.
+"""
+SQL Queries used during MariaDB ingestion
+"""
+
+MARIADB_GET_STORED_PROCEDURES = """
+SELECT
+    ROUTINE_NAME AS procedure_name,
+    ROUTINE_SCHEMA AS schema_name,
+    ROUTINE_DEFINITION AS definition,
+    ROUTINE_TYPE AS procedure_type,
+    ROUTINE_COMMENT AS description,
+    EXTERNAL_LANGUAGE AS language
+FROM
+    information_schema.ROUTINES
+WHERE
+    ROUTINE_SCHEMA = '{schema_name}'
+    AND ROUTINE_TYPE = 'PROCEDURE'
+"""
+
+MARIADB_GET_FUNCTIONS = """
+SELECT
+    ROUTINE_NAME AS procedure_name,
+    ROUTINE_SCHEMA AS schema_name,
+    ROUTINE_DEFINITION AS definition,
+    ROUTINE_TYPE AS procedure_type,
+    ROUTINE_COMMENT AS description,
+    EXTERNAL_LANGUAGE AS language
+FROM
+    information_schema.ROUTINES
+WHERE
+    ROUTINE_SCHEMA = '{schema_name}'
+    AND ROUTINE_TYPE = 'FUNCTION'
+"""

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

@@ -48,7 +48,9 @@ from metadata.ingestion.connections.test_connections import test_connection_db_c
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.oracle.queries import (
     CHECK_ACCESS_TO_ALL,
+    TEST_MATERIALIZED_VIEWS,
     TEST_ORACLE_GET_STORED_PACKAGES,
+    TEST_QUERY_HISTORY,
 )
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
@@ -99,6 +101,8 @@ class OracleConnection(BaseConnection[OracleConnectionConfig, Engine]):
         test_conn_queries = {
             "CheckAccess": CHECK_ACCESS_TO_ALL,
             "PackageAccess": TEST_ORACLE_GET_STORED_PACKAGES,
+            "GetMaterializedViews": TEST_MATERIALIZED_VIEWS,
+            "GetQueryHistory": TEST_QUERY_HISTORY,
         }
 
         return test_connection_db_common(

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

@@ -135,6 +135,24 @@ WHERE
 )
 
 CHECK_ACCESS_TO_ALL = "SELECT table_name FROM DBA_TABLES where ROWNUM < 2"
+
+
+TEST_MATERIALIZED_VIEWS = textwrap.dedent(
+    """
+SELECT COUNT(*) as count
+FROM DBA_MVIEWS
+WHERE ROWNUM = 1
+"""
+)
+
+TEST_QUERY_HISTORY = textwrap.dedent(
+    """
+SELECT COUNT(*) as count
+FROM gv$sql
+WHERE ROWNUM = 1
+"""
+)
+
 ORACLE_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent(
     """
 WITH SP_HISTORY AS (SELECT

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

@@ -43,6 +43,9 @@ from metadata.ingestion.source.database.postgres.queries import (
     POSTGRES_GET_DATABASE,
     POSTGRES_TEST_GET_QUERIES,
     POSTGRES_TEST_GET_TAGS,
+    TEST_COLUMN_METADATA,
+    TEST_INFORMATION_SCHEMA_COLUMNS,
+    TEST_TABLE_COMMENTS,
 )
 from metadata.ingestion.source.database.postgres.utils import (
     get_postgres_time_column_name,
@@ -91,6 +94,9 @@ class PostgresConnection(BaseConnection[PostgresConnectionConfig, Engine]):
             ),
             "GetDatabases": POSTGRES_GET_DATABASE,
             "GetTags": POSTGRES_TEST_GET_TAGS,
+            "GetColumnMetadata": TEST_COLUMN_METADATA,
+            "GetTableComments": TEST_TABLE_COMMENTS,
+            "GetInformationSchemaColumns": TEST_INFORMATION_SCHEMA_COLUMNS,
         }
         return test_connection_db_common(
             metadata=metadata,

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

@@ -249,3 +249,25 @@ WHERE
     prokind = 'f'
     and pg_namespace.nspname = '{schema_name}';
 """
+
+TEST_COLUMN_METADATA = """
+SELECT COUNT(*) as count
+FROM pg_catalog.pg_attribute a
+LEFT JOIN pg_catalog.pg_description pgd ON (
+    pgd.objoid = a.attrelid AND pgd.objsubid = a.attnum)
+WHERE 1=0
+"""
+
+TEST_TABLE_COMMENTS = """
+SELECT COUNT(*) as count
+FROM pg_catalog.pg_class c
+LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+LEFT JOIN pg_catalog.pg_description pgd ON pgd.objsubid = 0 AND pgd.objoid = c.oid
+WHERE 1=0
+"""
+
+TEST_INFORMATION_SCHEMA_COLUMNS = """
+SELECT COUNT(*) as count
+FROM information_schema.columns
+WHERE 1=0
+"""

metadata/ingestion/source/database/sample_data.py

@@ -151,6 +151,7 @@ from metadata.generated.schema.type.entityLineage import (
     LineageDetails,
 )
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.generated.schema.type.lifeCycle import AccessDetails, LifeCycle
 from metadata.generated.schema.type.schema import Topic as TopicSchema
 from metadata.ingestion.api.common import Entity
@@ -192,7 +193,7 @@ NUM_SERVICES = 1
 DATABASES_PER_SERVICE = 5
 SCHEMAS_PER_DATABASE = 5
 TABLES_PER_SCHEMA = 10
-COLUMNS_PER_TABLE = 50
+COLUMNS_PER_TABLE = 200
 NUM_THREADS = 10
 BATCH_SIZE = 10
 COLUMNS = [
@@ -1064,13 +1065,12 @@ class SampleDataSource(
             if directory_data.get("parent"):
                 parent_name = directory_data["parent"]
                 if parent_name in directory_refs:
-                    directory_request.parent = directory_refs[parent_name]
+                    directory_request.parent = FullyQualifiedEntityName(
+                        root=directory_refs[parent_name]
+                    )
                 else:
-                    # For nested references like "Marketing.Campaigns_2024"
-                    # Build parent FQN manually
-                    parent_path = parent_name.replace(".", "/")
-                    directory_request.parent = (
-                        f"{self.drive_service.fullyQualifiedName.root}.{parent_path}"
+                    directory_request.parent = FullyQualifiedEntityName(
+                        f"{self.drive_service.fullyQualifiedName.root}.{parent_name}"
                     )
 
             # Use direct API call instead of yielding since suffix mapping is missing
@@ -2491,9 +2491,25 @@ class SampleDataSource(
 
             # Create with minimal required fields
             try:
+                owner = self.metadata.get_by_name(User, "admin")
                 table_request = Either(
                     right=CreateTableRequest(
-                        name=table_name, databaseSchema=schema_fqn, columns=COLUMNS
+                        name=table_name,
+                        databaseSchema=schema_fqn,
+                        columns=COLUMNS,
+                        description=random.choice(
+                            [f"This is {table_name} description.", None]
+                        ),
+                        owners=random.choice(
+                            [
+                                EntityReferenceList(
+                                    [EntityReference(id=owner.id, type="user")]
+                                ),
+                                None,
+                            ]
+                        )
+                        if owner
+                        else None,
                     )
                 )
                 yield table_request

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

@@ -22,7 +22,6 @@ from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_SQL_STATEMENT,
 )
 from metadata.ingestion.source.database.snowflake.query_parser import (
-    SNOWFLAKE_QUERY_BATCH_SIZE,
     SnowflakeQueryParserSource,
 )
 from metadata.ingestion.source.database.stored_procedures_mixin import (
@@ -73,48 +72,26 @@ class SnowflakeLineageSource(
         yield a TableQuery with query parsing info
         """
         for engine in self.get_engine():
-            offset = 0
-            total_fetched = 0
-            max_results = self.source_config.resultLimit
-            while total_fetched < max_results:
-                batch_size = min(
-                    SNOWFLAKE_QUERY_BATCH_SIZE, max_results - total_fetched
+            rows = []
+            with engine.connect() as conn:
+                rows = conn.execution_options(
+                    stream_results=True, max_row_buffer=100
+                ).execute(
+                    self.get_sql_statement(start_time=self.start, end_time=self.end)
                 )
-                rows = []
-                row_count = 0
-                with engine.connect() as conn:
-                    rows = conn.execution_options(
-                        stream_results=True, max_row_buffer=100
-                    ).execute(
-                        self.get_sql_statement(
-                            start_time=self.start,
-                            end_time=self.end,
-                            offset=offset,
-                            limit=batch_size,
-                        )
+            # exit from active connection after fetching rows & during
+            # further process of `yield_query_lineage`
+            for row in rows:
+                query_dict = dict(row)
+                query_dict.update({k.lower(): v for k, v in query_dict.items()})
+                try:
+                    yield TableQuery(
+                        dialect=self.dialect.value,
+                        query=query_dict["query_text"],
+                        databaseName=self.get_database_name(query_dict),
+                        serviceName=self.config.serviceName,
+                        databaseSchema=self.get_schema_name(query_dict),
                     )
-                for row in rows:
-                    query_dict = dict(row)
-                    query_dict.update({k.lower(): v for k, v in query_dict.items()})
-                    row_count += 1
-                    try:
-                        yield TableQuery(
-                            dialect=self.dialect.value,
-                            query=query_dict["query_text"],
-                            databaseName=self.get_database_name(query_dict),
-                            serviceName=self.config.serviceName,
-                            databaseSchema=self.get_schema_name(query_dict),
-                        )
-                    except Exception as exc:
-                        logger.debug(traceback.format_exc())
-                        logger.warning(
-                            f"Error processing query_dict {query_dict}: {exc}"
-                        )
-                total_fetched += row_count
-                if row_count < batch_size:
-                    break
-                offset += batch_size
-                logger.info(
-                    f"Fetching next page with offset {offset} (fetched {total_fetched}/{max_results}) "
-                    f"for lineage queries"
-                )
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(f"Error processing query_dict {query_dict}: {exc}")

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

@@ -31,9 +31,7 @@ SNOWFLAKE_SQL_STATEMENT = textwrap.dedent(
     AND query_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
     AND start_time between to_timestamp_ltz('{start_time}') and to_timestamp_ltz('{end_time}')
     {filters}
-    ORDER BY start_time
     LIMIT {result_limit}
-    OFFSET {offset}
     """
 )
 

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

@@ -32,7 +32,6 @@ from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 SNOWFLAKE_ABORTED_CODE = "1969"
-SNOWFLAKE_QUERY_BATCH_SIZE = 1000
 
 
 class SnowflakeQueryParserSource(QueryParserSource, ABC):
@@ -52,27 +51,18 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
             )
         return cls(config, metadata)
 
-    def get_sql_statement(
-        self,
-        start_time: datetime,
-        end_time: datetime,
-        offset: int = 0,
-        limit: int = None,
-    ) -> str:
+    def get_sql_statement(self, start_time: datetime, end_time: datetime) -> str:
         """
         returns sql statement to fetch query logs
         """
-        if limit is None:
-            limit = self.config.sourceConfig.config.resultLimit
         return self.sql_stmt.format(
             start_time=start_time,
             end_time=end_time,
-            result_limit=limit,
+            result_limit=self.config.sourceConfig.config.resultLimit,
             filters=self.get_filters(),
             account_usage=self.service_connection.accountUsageSchema,
             credit_cost=self.service_connection.creditCost
             * self.service_connection.creditCost,
-            offset=offset,
         )
 
     def check_life_cycle_query(

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

@@ -11,22 +11,11 @@
 """
 Snowflake usage module
 """
-import traceback
-from datetime import timedelta
-from typing import Iterable
-
-from metadata.generated.schema.type.basic import DateTime
-from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
-from metadata.ingestion.lineage.masker import mask_query
 from metadata.ingestion.source.database.snowflake.queries import SNOWFLAKE_SQL_STATEMENT
 from metadata.ingestion.source.database.snowflake.query_parser import (
-    SNOWFLAKE_QUERY_BATCH_SIZE,
     SnowflakeQueryParserSource,
 )
 from metadata.ingestion.source.database.usage_source import UsageSource
-from metadata.utils.logger import ingestion_logger
-
-logger = ingestion_logger()
 
 
 class SnowflakeUsageSource(SnowflakeQueryParserSource, UsageSource):
@@ -52,89 +41,3 @@ class SnowflakeUsageSource(SnowflakeQueryParserSource, UsageSource):
         "INSERT",
         "MERGE",
     ]
-
-    def yield_table_queries(self) -> Iterable[TableQuery]:
-        """
-        Given an Engine, iterate over the day range and
-        query the results with pagination
-        """
-        daydiff = self.end - self.start
-        for days in range(daydiff.days):
-            logger.info(
-                f"Scanning query logs for {(self.start + timedelta(days=days)).date()} - "
-                f"{(self.start + timedelta(days=days + 1)).date()}"
-            )
-            query = None
-            offset = 0
-            total_fetched = 0
-            max_results = self.source_config.resultLimit
-            try:
-                for engine in self.get_engine():
-                    while total_fetched < max_results:
-                        batch_size = min(
-                            SNOWFLAKE_QUERY_BATCH_SIZE, max_results - total_fetched
-                        )
-                        query = self.get_sql_statement(
-                            start_time=self.start + timedelta(days=days),
-                            end_time=self.start + timedelta(days=days + 1),
-                            offset=offset,
-                            limit=batch_size,
-                        )
-                        with engine.connect() as conn:
-                            rows = conn.execute(query)
-                            queries = []
-                            row_count = 0
-                            for row in rows:
-                                row = dict(row)
-                                row_count += 1
-                                try:
-                                    row.update({k.lower(): v for k, v in row.items()})
-                                    logger.debug(f"Processing row: {row}")
-                                    query_type = row.get("query_type")
-                                    query_text = self.format_query(row["query_text"])
-                                    queries.append(
-                                        TableQuery(
-                                            query=query_text,
-                                            query_type=query_type,
-                                            exclude_usage=self.check_life_cycle_query(
-                                                query_type=query_type,
-                                                query_text=query_text,
-                                            ),
-                                            dialect=self.dialect.value,
-                                            userName=row["user_name"],
-                                            startTime=str(row["start_time"]),
-                                            endTime=str(row["end_time"]),
-                                            analysisDate=DateTime(row["start_time"]),
-                                            aborted=self.get_aborted_status(row),
-                                            databaseName=self.get_database_name(row),
-                                            duration=row.get("duration"),
-                                            serviceName=self.config.serviceName,
-                                            databaseSchema=self.get_schema_name(row),
-                                            cost=row.get("cost"),
-                                        )
-                                    )
-                                except Exception as exc:
-                                    logger.debug(traceback.format_exc())
-                                    logger.warning(
-                                        f"Unexpected exception processing row [{row}]: {exc}"
-                                    )
-                        if queries:
-                            yield TableQueries(queries=queries)
-                        total_fetched += row_count
-                        if row_count < batch_size:
-                            break
-                        offset += batch_size
-                        logger.info(
-                            f"Fetching next page with offset {offset} (fetched {total_fetched}/{max_results}) "
-                            f"for {(self.start + timedelta(days=days)).date()}"
-                        )
-            except Exception as exc:
-                if query:
-                    logger.debug(
-                        (
-                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value) or query}"
-                            "\n##########################"
-                        )
-                    )
-                logger.debug(traceback.format_exc())
-                logger.error(f"Source usage processing error: {exc}")

metadata/ingestion/source/database/timescale/__init__.py

@@ -0,0 +1,3 @@
+from metadata.ingestion.source.database.timescale.metadata import TimescaleSource
+
+__all__ = ["TimescaleSource"]