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

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}")

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

@@ -0,0 +1,228 @@
+#  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 source module
+"""
+import traceback
+from typing import Iterable, Optional, Tuple
+
+from sqlalchemy.engine import Inspector
+
+from metadata.generated.schema.entity.data.table import (
+    PartitionColumnDetails,
+    PartitionIntervalTypes,
+    Table,
+    TablePartition,
+)
+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.models import Either
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.postgres.metadata import PostgresSource
+from metadata.ingestion.source.database.timescale.models import (
+    CompressionSettings,
+    HypertableInfo,
+)
+from metadata.ingestion.source.database.timescale.queries import (
+    TIMESCALE_CHECK_EXTENSION,
+    TIMESCALE_GET_COMPRESSION_SETTINGS,
+    TIMESCALE_GET_HYPERTABLE_INFO,
+)
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class TimescaleSource(PostgresSource):
+    """
+    Implements the necessary methods to extract
+    Database metadata from TimescaleDB Source
+    """
+
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+        self.timescaledb_installed = False
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: TimescaleConnection = config.serviceConnection.root.config
+        if not isinstance(connection, TimescaleConnection):
+            raise InvalidSourceException(
+                f"Expected TimescaleConnection, but got {connection}"
+            )
+        return cls(config, metadata)
+
+    def prepare(self):
+        """
+        Check if TimescaleDB extension is installed
+        """
+        super().prepare()
+        try:
+            result = self.engine.execute(TIMESCALE_CHECK_EXTENSION).first()
+            if result:
+                self.timescaledb_installed = result.timescaledb_installed
+                logger.info(
+                    f"TimescaleDB extension installed: {self.timescaledb_installed}"
+                )
+        except Exception as exc:
+            logger.warning(f"Could not check TimescaleDB extension: {exc}")
+            self.timescaledb_installed = False
+
+    def yield_table(
+        self, table_name_and_type: Tuple[str, str]
+    ) -> Iterable[Either[Table]]:
+        """
+        Override to add TimescaleDB-specific metadata
+        """
+        for either_table in super().yield_table(table_name_and_type):
+            if either_table.right and self.timescaledb_installed:
+                try:
+                    self._add_timescale_metadata(either_table.right)
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(
+                        f"Error adding TimescaleDB metadata for table {either_table.right.name}: {exc}"
+                    )
+            yield either_table
+
+    def _add_timescale_metadata(self, table: Table) -> None:
+        """
+        Add compression and hypertable metadata to the table
+        """
+        try:
+            schema_name = table.databaseSchema.root
+            table_name = table.name.root
+
+            hypertable = self._get_hypertable_info(table_name, schema_name)
+
+            if hypertable:
+                logger.debug(f"Found hypertable: {schema_name}.{table_name}")
+
+                table.compressionEnabled = hypertable.compression_enabled
+
+                if hypertable.compression_enabled:
+                    compression = self._get_compression_settings(
+                        table_name, schema_name
+                    )
+
+                    if compression:
+                        table.compressionCodec = "TimescaleDB Native"
+                        table.compressionStrategy = {
+                            "compressionType": "POLICY_BASED",
+                            "segmentColumns": compression.segment_by_columns or [],
+                            "orderColumns": compression.order_by_columns or [],
+                        }
+
+                partition_info = self._build_hypertable_partition(hypertable)
+                if partition_info:
+                    table.tablePartition = partition_info
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error processing TimescaleDB metadata for {table.name}: {exc}"
+            )
+
+    def _get_hypertable_info(
+        self, table_name: str, schema_name: str
+    ) -> Optional[HypertableInfo]:
+        """
+        Query timescaledb_information.hypertables for metadata
+        """
+        try:
+            result = self.engine.execute(
+                TIMESCALE_GET_HYPERTABLE_INFO,
+                {"table_name": table_name, "schema_name": schema_name},
+            ).first()
+
+            if result:
+                return HypertableInfo.model_validate(dict(result._mapping))
+            return None
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.debug(
+                f"Could not get hypertable info for {schema_name}.{table_name}: {exc}"
+            )
+            return None
+
+    def _get_compression_settings(
+        self, table_name: str, schema_name: str
+    ) -> Optional[CompressionSettings]:
+        """
+        Query timescaledb_information.compression_settings for compression config
+        """
+        try:
+            result = self.engine.execute(
+                TIMESCALE_GET_COMPRESSION_SETTINGS,
+                {"table_name": table_name, "schema_name": schema_name},
+            ).first()
+
+            if result:
+                return CompressionSettings.model_validate(dict(result._mapping))
+            return None
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.debug(
+                f"Could not get compression settings for {schema_name}.{table_name}: {exc}"
+            )
+            return None
+
+    def _build_hypertable_partition(
+        self, hypertable: HypertableInfo
+    ) -> Optional[TablePartition]:
+        """
+        Build partition details from hypertable information
+        """
+        if not hypertable.column_name:
+            return None
+
+        if hypertable.interval_length:
+            interval_type = PartitionIntervalTypes.TIME_UNIT
+            interval = f"{hypertable.interval_length} microseconds"
+        elif hypertable.integer_interval:
+            interval_type = PartitionIntervalTypes.INTEGER_RANGE
+            interval = str(hypertable.integer_interval)
+        else:
+            interval_type = PartitionIntervalTypes.COLUMN_VALUE
+            interval = None
+
+        return TablePartition(
+            columns=[
+                PartitionColumnDetails(
+                    columnName=hypertable.column_name,
+                    intervalType=interval_type,
+                    interval=interval,
+                )
+            ]
+        )
+
+    def get_table_partition_details(
+        self, table_name: str, schema_name: str, inspector: Inspector
+    ) -> Tuple[bool, Optional[TablePartition]]:
+        """
+        Override to check for hypertables first, then fall back to PostgreSQL partitioning
+        """
+        if self.timescaledb_installed:
+            hypertable = self._get_hypertable_info(table_name, schema_name)
+            if hypertable:
+                partition = self._build_hypertable_partition(hypertable)
+                return (True, partition) if partition else (False, None)
+
+        return super().get_table_partition_details(table_name, schema_name, inspector)

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

@@ -0,0 +1,60 @@
+#  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 models for metadata extraction
+"""
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class HypertableInfo(BaseModel):
+    """Information about a TimescaleDB hypertable"""
+
+    hypertable_schema: str
+    hypertable_name: str
+    compression_enabled: bool
+    column_name: Optional[str] = None
+    interval_length: Optional[int] = None
+    integer_interval: Optional[int] = None
+    integer_now_func: Optional[str] = None
+    num_dimensions: int = 1
+
+
+class CompressionSettings(BaseModel):
+    """Compression configuration for a hypertable"""
+
+    segment_by_columns: Optional[List[str]] = Field(default_factory=list)
+    order_by_columns: Optional[List[str]] = Field(default_factory=list)
+
+
+class ContinuousAggregateInfo(BaseModel):
+    """Information about a continuous aggregate"""
+
+    view_schema: str
+    view_name: str
+    view_definition: Optional[str] = None
+    compression_enabled: bool = False
+    materialized_only: bool = False
+    materialization_hypertable_schema: Optional[str] = None
+    materialization_hypertable_name: Optional[str] = None
+
+
+class ChunkInfo(BaseModel):
+    """Information about a hypertable chunk"""
+
+    chunk_schema: str
+    chunk_name: str
+    range_start: Optional[str] = None
+    range_end: Optional[str] = None
+    is_compressed: bool = False
+    chunk_tablespace: Optional[str] = None

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

@@ -0,0 +1,96 @@
+#  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 system view queries
+"""
+
+TIMESCALE_GET_HYPERTABLE_INFO = """
+SELECT
+    ht.hypertable_schema,
+    ht.hypertable_name,
+    ht.compression_enabled,
+    dim.column_name,
+    EXTRACT(EPOCH FROM dim.time_interval) * 1000000 as interval_length,
+    dim.integer_interval,
+    dim.integer_now_func,
+    ht.num_dimensions
+FROM timescaledb_information.hypertables ht
+LEFT JOIN timescaledb_information.dimensions dim
+    ON ht.hypertable_schema = dim.hypertable_schema
+    AND ht.hypertable_name = dim.hypertable_name
+    AND dim.dimension_number = 1
+WHERE ht.hypertable_schema = %(schema_name)s
+  AND ht.hypertable_name = %(table_name)s
+"""
+
+TIMESCALE_GET_COMPRESSION_SETTINGS = """
+SELECT
+    array_agg(DISTINCT attname ORDER BY attname) FILTER (WHERE segmentby_column_index IS NOT NULL) as segment_by_columns,
+    array_agg(DISTINCT attname ORDER BY attname) FILTER (WHERE orderby_column_index IS NOT NULL) as order_by_columns
+FROM timescaledb_information.compression_settings
+WHERE hypertable_schema = %(schema_name)s
+  AND hypertable_name = %(table_name)s
+GROUP BY hypertable_schema, hypertable_name
+"""
+
+TIMESCALE_GET_CONTINUOUS_AGGREGATES = """
+SELECT
+    view_schema,
+    view_name,
+    view_definition,
+    compression_enabled,
+    materialized_only,
+    materialization_hypertable_schema,
+    materialization_hypertable_name
+FROM timescaledb_information.continuous_aggregates
+WHERE view_schema = %(schema_name)s
+"""
+
+TIMESCALE_GET_CHUNK_INFO = """
+SELECT
+    chunk_schema,
+    chunk_name,
+    range_start,
+    range_end,
+    is_compressed,
+    chunk_tablespace
+FROM timescaledb_information.chunks
+WHERE hypertable_schema = %(schema_name)s
+  AND hypertable_name = %(table_name)s
+ORDER BY range_start DESC
+LIMIT 5
+"""
+
+TIMESCALE_CHECK_EXTENSION = """
+SELECT EXISTS (
+    SELECT 1
+    FROM pg_extension
+    WHERE extname = 'timescaledb'
+) as timescaledb_installed
+"""
+
+TIMESCALE_GET_CONTINUOUS_AGGREGATE_DEFINITIONS = """
+SELECT
+    ca.view_schema,
+    ca.view_name,
+    ca.view_definition,
+    ca.materialization_hypertable_schema,
+    ca.materialization_hypertable_name,
+    ht.hypertable_schema as source_hypertable_schema,
+    ht.hypertable_name as source_hypertable_name
+FROM timescaledb_information.continuous_aggregates ca
+LEFT JOIN timescaledb_information.hypertables ht
+    ON ca.materialization_hypertable_schema = ht.hypertable_schema
+    AND ca.materialization_hypertable_name = ht.hypertable_name
+WHERE ca.view_schema NOT IN ('_timescaledb_catalog', '_timescaledb_config', '_timescaledb_cache', '_timescaledb_internal')
+ORDER BY ca.view_schema, ca.view_name
+"""