openmetadata-ingestion 1.9.11.2__py3-none-any.whl → 1.10.0.0__py3-none-any.whl

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/constants.py

@@ -14,6 +14,9 @@ Snowflake constants
 
 from sqlalchemy.sql.sqltypes import BOOLEANTYPE, VARCHAR
 
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureType
+from metadata.generated.schema.entity.data.table import TableType
+
 DEFAULT_STREAM_COLUMNS = [
     {
         "name": "METADATA$ACTION",
@@ -49,3 +52,18 @@ DEFAULT_STREAM_COLUMNS = [
 
 SNOWFLAKE_TAG_DESCRIPTION = "SNOWFLAKE TAG VALUE"
 SNOWFLAKE_CLASSIFICATION_DESCRIPTION = "SNOWFLAKE TAG NAME"
+
+TABLE_TYPE_URL_MAP = {
+    TableType.Regular: "table",
+    TableType.Transient: "table",
+    TableType.Dynamic: "dynamic-table",
+    TableType.External: "external-table",
+    TableType.View: "view",
+    TableType.MaterializedView: "materialized-view",
+    TableType.Stream: "stream",
+}
+
+PROCEDURE_TYPE_URL_MAP = {
+    StoredProcedureType.StoredProcedure: "procedure",
+    StoredProcedureType.UDF: "user-function",
+}

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

@@ -20,6 +20,7 @@ import sqlalchemy.types as sqltypes
 import sqlparse
 from snowflake.sqlalchemy.custom_types import VARIANT, StructuredType
 from snowflake.sqlalchemy.snowdialect import SnowflakeDialect, ischema_names
+from sqlalchemy import exc as sa_exc
 from sqlalchemy.engine.reflection import Inspector
 from sqlparse.sql import Function, Identifier, Token
 
@@ -74,8 +75,10 @@ from metadata.ingestion.source.database.incremental_metadata_extraction import (
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.snowflake.constants import (
     DEFAULT_STREAM_COLUMNS,
+    PROCEDURE_TYPE_URL_MAP,
     SNOWFLAKE_CLASSIFICATION_DESCRIPTION,
     SNOWFLAKE_TAG_DESCRIPTION,
+    TABLE_TYPE_URL_MAP,
 )
 from metadata.ingestion.source.database.snowflake.models import (
     STORED_PROC_LANGUAGE_MAP,
@@ -688,11 +691,11 @@ class SnowflakeSource(
         table_type: Optional[TableType] = None,
     ) -> Optional[str]:
         """
-        Method to get the source url for snowflake
+        Method to get the source url for snowflake tables
         """
         try:
             if self.account and self.org_name:
-                tab_type = "view" if table_type == TableType.View else "table"
+                tab_type = TABLE_TYPE_URL_MAP.get(table_type, "table")
                 url = self._get_source_url_root(
                     database_name=database_name, schema_name=schema_name
                 )
@@ -704,6 +707,41 @@ class SnowflakeSource(
             logger.error(f"Unable to get source url: {exc}")
         return None
 
+    def get_procedure_source_url(
+        self,
+        database_name: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        procedure_name: Optional[str] = None,
+        procedure_signature: Optional[str] = None,
+        procedure_type: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Method to get the source url for snowflake stored procedures
+        """
+        try:
+            if self.account and self.org_name:
+                url = self._get_source_url_root(
+                    database_name=database_name, schema_name=schema_name
+                )
+
+                # Convert string procedure type to enum and get URL mapping
+                proc_type_enum = (
+                    StoredProcedureType(procedure_type)
+                    if procedure_type
+                    else StoredProcedureType.StoredProcedure
+                )
+                tab_type = PROCEDURE_TYPE_URL_MAP.get(proc_type_enum, "procedure")
+
+                if procedure_name:
+                    full_name = f"{procedure_name}{procedure_signature or ''}"
+                    url = f"{url}/{tab_type}/{full_name}"
+
+                return url
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(f"Unable to get procedure source url: {exc}")
+        return None
+
     def _get_view_names_and_types(
         self, schema_name: str, materialized_views: bool = False
     ) -> List[TableNameAndType]:
@@ -844,12 +882,13 @@ class SnowflakeSource(
                     schema_name=self.context.get().database_schema,
                 ),
                 sourceUrl=SourceUrl(
-                    self._get_source_url_root(
+                    self.get_procedure_source_url(
                         database_name=self.context.get().database,
                         schema_name=self.context.get().database_schema,
+                        procedure_name=stored_procedure.name,
+                        procedure_signature=stored_procedure.signature,
+                        procedure_type=stored_procedure.procedure_type,
                     )
-                    + f"/procedure/{stored_procedure.name}"
-                    + f"{stored_procedure.signature if stored_procedure.signature else ''}"
                 ),
             )
             yield Either(right=stored_procedure_request)
@@ -910,17 +949,26 @@ class SnowflakeSource(
                     table_name = result[6].split(".")[-1]
                     # Can't fetch source of stream is source is dropped or no priviledge
                     if table_name == "No privilege or table dropped":
-                        logger.debug(
-                            f"Couldn't fetch columns of stream [{result and result[1]}],"
-                            f" due to error on source: {table_name}. Result: {result}"
+                        logger.warning(
+                            f"Couldn't fetch columns of stream [{result and result[1]}] "
+                            f"(schema: '{schema_name}', db: '{db_name}') due to error on"
+                            f" source: [{table_name}]. Result: {result}"
                         )
                         return []
             except Exception:
                 pass
 
-        columns = inspector.get_columns(
-            table_name, schema_name, table_type=table_type, db_name=db_name
-        )
+        try:
+            columns = inspector.get_columns(
+                table_name, schema_name, table_type=table_type, db_name=db_name
+            )
+        except sa_exc.NoSuchTableError:
+            logger.warning(
+                f"Table [{table_name}] (schema: '{schema_name}', db: '{db_name}') not found."
+                " Unable to fetch columns. Please check if the configured Snowflake user has"
+                " necessary grants on this table."
+            )
+            return []
 
         if table_type == TableType.Stream:
             columns = [*columns, *DEFAULT_STREAM_COLUMNS]

metadata/ingestion/source/database/sql_column_handler.py

@@ -268,9 +268,17 @@ class SqlColumnHandlerMixin:
 
         table_columns = []
 
-        columns = self._get_columns_internal(
-            schema_name, table_name, db_name, inspector, table_type
-        )
+        try:
+            columns = self._get_columns_internal(
+                schema_name, table_name, db_name, inspector, table_type
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unexpected exception getting columns for table [{table_name}] "
+                f"(schema: '{schema_name}', db: '{db_name}'): {exc}"
+            )
+            columns = []
 
         def process_column(column: dict):
             (

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

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

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

@@ -0,0 +1,103 @@
+#  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.
+
+"""
+Source connection handler for TimescaleDB
+"""
+from typing import Optional
+
+from sqlalchemy.engine import Engine
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.common.azureConfig import (
+    AzureConfigurationSource,
+)
+from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
+    BasicAuth,
+)
+from metadata.generated.schema.entity.services.connections.database.timescaleConnection import (
+    TimescaleConnection as TimescaleConnectionConfig,
+)
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
+from metadata.ingestion.connections.builders import (
+    create_generic_db_connection,
+    get_connection_args_common,
+    get_connection_url_common,
+)
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.connections.test_connections import test_connection_db_common
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.postgres.queries import (
+    POSTGRES_GET_DATABASE,
+    POSTGRES_TEST_GET_QUERIES,
+    POSTGRES_TEST_GET_TAGS,
+)
+from metadata.ingestion.source.database.postgres.utils import (
+    get_postgres_time_column_name,
+)
+from metadata.utils.constants import THREE_MIN
+from metadata.utils.credentials import get_azure_access_token
+
+
+class TimescaleConnection(BaseConnection[TimescaleConnectionConfig, Engine]):
+    def _get_client(self) -> Engine:
+        """
+        Return the SQLAlchemy Engine for TimescaleDB.
+        TimescaleDB is a PostgreSQL extension, so we use the same connection logic.
+        """
+        connection = self.service_connection
+
+        if isinstance(connection.authType, AzureConfigurationSource):
+            access_token = get_azure_access_token(connection.authType)
+            connection.authType = BasicAuth(password=access_token)
+        return create_generic_db_connection(
+            connection=connection,
+            get_connection_url_fn=get_connection_url_common,
+            get_connection_args_fn=get_connection_args_common,
+        )
+
+    def get_connection_dict(self) -> dict:
+        """
+        Return the connection dictionary for this service.
+        """
+        raise NotImplementedError(
+            "get_connection_dict is not implemented for TimescaleDB"
+        )
+
+    def test_connection(
+        self,
+        metadata: OpenMetadata,
+        automation_workflow: Optional[AutomationWorkflow] = None,
+        timeout_seconds: Optional[int] = THREE_MIN,
+    ) -> TestConnectionResult:
+        """
+        Test connection. This can be executed either as part
+        of a metadata workflow or during an Automation Workflow
+        """
+        queries = {
+            "GetQueries": POSTGRES_TEST_GET_QUERIES.format(
+                time_column_name=get_postgres_time_column_name(engine=self.client),
+            ),
+            "GetDatabases": POSTGRES_GET_DATABASE,
+            "GetTags": POSTGRES_TEST_GET_TAGS,
+        }
+        return test_connection_db_common(
+            metadata=metadata,
+            engine=self.client,
+            service_connection=self.service_connection,
+            automation_workflow=automation_workflow,
+            timeout_seconds=timeout_seconds,
+            queries=queries,
+        )

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

@@ -0,0 +1,124 @@
+#  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.
+
+"""
+TimescaleDB lineage module with continuous aggregate support
+"""
+
+import traceback
+from typing import Iterable, Optional
+
+from metadata.generated.schema.entity.services.connections.database.timescaleConnection import (
+    TimescaleConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.postgres.lineage import PostgresLineageSource
+from metadata.ingestion.source.database.timescale.queries import (
+    TIMESCALE_CHECK_EXTENSION,
+    TIMESCALE_GET_CONTINUOUS_AGGREGATE_DEFINITIONS,
+)
+from metadata.ingestion.source.models import TableView
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class TimescaleLineageSource(PostgresLineageSource):
+    """
+    TimescaleDB lineage source that extends PostgreSQL lineage
+    to include continuous aggregates (auto-refreshing materialized views)
+    """
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        """Create TimescaleLineageSource"""
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection = config.serviceConnection.root.config
+        if not isinstance(connection, (TimescaleConnection)):
+            raise InvalidSourceException(
+                f"Expected TimescaleConnection, but got {connection}"
+            )
+        return cls(config, metadata)
+
+    def __init__(self, config, metadata):
+        super().__init__(config, metadata)
+        self.timescaledb_installed = False
+
+    def prepare(self):
+        """Check if TimescaleDB extension is installed"""
+        super().prepare()
+        try:
+            result = self.engine.execute(TIMESCALE_CHECK_EXTENSION).first()
+            self.timescaledb_installed = (
+                result.timescaledb_installed if result else False
+            )
+            logger.info(
+                f"TimescaleDB extension installed for lineage: {self.timescaledb_installed}"
+            )
+        except Exception as exc:
+            logger.warning(f"Could not check TimescaleDB extension: {exc}")
+            self.timescaledb_installed = False
+
+    def view_lineage_producer(self) -> Iterable[TableView]:
+        """
+        Get view definitions from ES and add continuous aggregates
+        """
+        # First yield regular views from parent
+        yield from super().view_lineage_producer()
+
+        # Then yield continuous aggregates if TimescaleDB is installed
+        if self.timescaledb_installed:
+            yield from self._yield_continuous_aggregate_views()
+
+    def _yield_continuous_aggregate_views(self) -> Iterable[TableView]:
+        """
+        Get TimescaleDB continuous aggregates as TableView objects.
+
+        Continuous aggregates are auto-refreshing materialized views that
+        provide clear column-level lineage from source hypertables.
+
+        The view_definition contains the SELECT statement that defines
+        the continuous aggregate, which OpenMetadata can parse for
+        column-level lineage.
+        """
+        try:
+            results = self.engine.execute(
+                TIMESCALE_GET_CONTINUOUS_AGGREGATE_DEFINITIONS
+            )
+
+            for row in results:
+                try:
+                    yield TableView(
+                        table_name=row.view_name,
+                        schema_name=row.view_schema,
+                        db_name=self.service_connection.database,
+                        view_definition=row.view_definition,
+                    )
+
+                    logger.debug(
+                        f"Extracted continuous aggregate view: {row.view_schema}.{row.view_name}"
+                    )
+
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(
+                        f"Failed to process continuous aggregate {row.view_schema}.{row.view_name}: {exc}"
+                    )
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to extract continuous aggregate views: {exc}")