acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/datahub/datahub_database_reader.py

@@ -1,10 +1,10 @@
-import contextlib
 import json
 import logging
+import time
 from datetime import datetime
 from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, TypeVar
 
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, text
 
 from datahub.emitter.aspect import ASPECT_MAP
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -12,13 +12,14 @@ from datahub.emitter.serialization_helper import post_json_transform
 from datahub.ingestion.source.datahub.config import DataHubSourceConfig
 from datahub.ingestion.source.datahub.report import DataHubSourceReport
 from datahub.ingestion.source.sql.sql_config import SQLAlchemyConnectionConfig
-from datahub.metadata.schema_classes import ChangeTypeClass, SystemMetadataClass
+from datahub.metadata.schema_classes import SystemMetadataClass
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 
 logger = logging.getLogger(__name__)
 
 # Should work for at least mysql, mariadb, postgres
 DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S.%f"
+DATE_FORMAT = "%Y-%m-%d"
 
 ROW = TypeVar("ROW", bound=Dict[str, Any])
 
@@ -85,6 +86,9 @@ class DataHubDatabaseReader:
             **connection_config.options,
         )
 
+        # Cache for available dates to avoid redundant queries
+        self.available_dates_cache: Optional[List[datetime]] = None
+
     @property
     def soft_deleted_urns_query(self) -> str:
         return f"""
@@ -100,14 +104,28 @@ class DataHubDatabaseReader:
             ORDER BY mav.urn
         """
 
-    @property
-    def query(self) -> str:
-        # May repeat rows for the same date
-        # Offset is generally 0, unless we repeat the same createdon twice
+    def _get_json_extract_expression(self) -> str:
+        """
+        Returns the appropriate JSON extraction expression based on the database dialect.
+
+        Returns:
+            Database-specific JSON extraction expression
+        """
+        # Return the correct JSON extraction expression for the "removed" field,
+        # depending on the database dialect.
+        if self.engine.dialect.name == "postgresql":
+            # For PostgreSQL, cast the metadata column to JSON and extract the 'removed' key as boolean.
+            return "((metadata::json)->>'removed')::boolean"
+        else:
+            # For other databases (e.g., MySQL), use JSON_EXTRACT.
+            return "JSON_EXTRACT(metadata, '$.removed')"
+
+    def query(self, set_structured_properties_filter: bool) -> str:
+        """
+        Main query that gets data for specified date range with appropriate filters.
+        """
+        structured_prop_filter = f" AND urn {'' if set_structured_properties_filter else 'NOT'} like 'urn:li:structuredProperty:%%'"
 
-        # Ensures stable order, chronological per (urn, aspect)
-        # Relies on createdon order to reflect version order
-        # Ordering of entries with the same createdon is handled by VersionOrderer
         return f"""
         SELECT *
         FROM (
@@ -123,7 +141,7 @@ class DataHubDatabaseReader:
             LEFT JOIN (
                 SELECT
                     *,
-                    JSON_EXTRACT(metadata, '$.removed') as removed
+                    {self._get_json_extract_expression()} as removed
                 FROM {self.engine.dialect.identifier_preparer.quote(self.config.database_table_name)}
                 WHERE aspect = 'status'
                 AND version = 0
@@ -132,6 +150,7 @@ class DataHubDatabaseReader:
                 {"" if self.config.include_all_versions else "AND mav.version = 0"}
                 {"" if not self.config.exclude_aspects else "AND mav.aspect NOT IN %(exclude_aspects)s"}
                 AND mav.createdon >= %(since_createdon)s
+                AND mav.createdon < %(end_createdon)s
             ORDER BY
                 createdon,
                 urn,
@@ -139,50 +158,189 @@ class DataHubDatabaseReader:
                 version
         ) as t
         WHERE 1=1
-            {"" if self.config.include_soft_deleted_entities else "AND (removed = false or removed is NULL)"}
+            {"" if self.config.include_soft_deleted_entities else " AND (removed = false or removed is NULL)"}
+            {structured_prop_filter}
         ORDER BY
             createdon,
             urn,
             aspect,
             version
+        LIMIT %(limit)s
+        OFFSET %(offset)s
         """
 
+    def execute_with_params(
+        self, query: str, params: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+        """Execute query with proper parameter binding that works with your database"""
+        with self.engine.connect() as conn:
+            result = conn.execute(query, params or {})
+            return [dict(row) for row in result.fetchall()]
+
     def execute_server_cursor(
         self, query: str, params: Dict[str, Any]
     ) -> Iterable[Dict[str, Any]]:
+        """Execute a query with server-side cursor"""
         with self.engine.connect() as conn:
             if self.engine.dialect.name in ["postgresql", "mysql", "mariadb"]:
                 with (
                     conn.begin()
                 ):  # Transaction required for PostgreSQL server-side cursor
-                    # Note that stream_results=True is mainly supported by PostgreSQL and MySQL-based dialects.
-                    # https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.Connection.execution_options.params.stream_results
+                    # Set query timeout at the connection level
+                    if self.config.query_timeout:
+                        if self.engine.dialect.name == "postgresql":
+                            conn.execute(
+                                text(
+                                    f"SET statement_timeout = {self.config.query_timeout * 1000}"
+                                )
+                            )  # milliseconds
+                        elif self.engine.dialect.name in ["mysql", "mariadb"]:
+                            conn.execute(
+                                text(
+                                    f"SET max_execution_time = {self.config.query_timeout * 1000}"
+                                )
+                            )  # milliseconds
+
+                    # Stream results with batch size
                     conn = conn.execution_options(
                         stream_results=True,
                         yield_per=self.config.database_query_batch_size,
                     )
+
+                    # Execute query - using native parameterization without text()
+                    # to maintain compatibility with your original code
                     result = conn.execute(query, params)
                     for row in result:
                         yield dict(row)
+
+                return  # Success, exit the retry loop
             else:
                 raise ValueError(f"Unsupported dialect: {self.engine.dialect.name}")
 
     def _get_rows(
-        self, from_createdon: datetime, stop_time: datetime
+        self,
+        start_date: datetime,
+        end_date: datetime,
+        set_structured_properties_filter: bool,
+        limit: int,
     ) -> Iterable[Dict[str, Any]]:
-        params = {
-            "exclude_aspects": list(self.config.exclude_aspects),
-            "since_createdon": from_createdon.strftime(DATETIME_FORMAT),
-        }
-        yield from self.execute_server_cursor(self.query, params)
+        """
+        Retrieves data rows within a specified date range using pagination.
 
-    def get_aspects(
+        Implements a hybrid pagination strategy that switches between time-based and
+        offset-based approaches depending on the returned data. Uses server-side
+        cursors for efficient memory usage.
+
+        Note: May return duplicate rows across batch boundaries when multiple rows
+        share the same 'createdon' timestamp. This is expected behavior when
+        transitioning between pagination methods.
+
+        Args:
+        start_date: Beginning of date range (inclusive)
+        end_date: End of date range (exclusive)
+        set_structured_properties_filter: Whether to apply structured filtering
+        limit: Maximum rows to fetch per query
+
+        Returns:
+            An iterable of database rows as dictionaries
+        """
+        offset = 0
+        last_createdon = None
+        first_iteration = True
+
+        while True:
+            try:
+                # Set up query and parameters - using named parameters
+                query = self.query(set_structured_properties_filter)
+                params: Dict[str, Any] = {
+                    "since_createdon": start_date.strftime(DATETIME_FORMAT),
+                    "end_createdon": end_date.strftime(DATETIME_FORMAT),
+                    "limit": limit,
+                    "offset": offset,
+                    # Always pass exclude_aspects as a tuple, postgres doesn't support lists
+                    "exclude_aspects": tuple(self.config.exclude_aspects),
+                }
+
+                logger.info(
+                    f"Querying data from {start_date.strftime(DATETIME_FORMAT)} to {end_date.strftime(DATETIME_FORMAT)} "
+                    f"with limit {limit} and offset {offset} (inclusive range)"
+                )
+
+                # Execute query with server-side cursor
+                rows = self.execute_server_cursor(query, params)
+                # Process and yield rows
+                rows_processed = 0
+                for row in rows:
+                    if first_iteration:
+                        start_date = row.get("createdon", start_date)
+                        first_iteration = False
+
+                    last_createdon = row.get("createdon")
+                    rows_processed += 1
+                    yield row
+
+                # If we processed fewer than the limit or no last_createdon, we're done
+                if rows_processed < limit or not last_createdon:
+                    break
+
+                # Update parameters for next iteration
+                if start_date != last_createdon:
+                    start_date = last_createdon
+                    offset = 0
+                else:
+                    offset += limit
+
+                logger.info(
+                    f"Processed {rows_processed} rows for date range {start_date} to {end_date}. Continuing to next batch."
+                )
+
+            except Exception as e:
+                logger.error(
+                    f"Error processing date range {start_date} to {end_date}: {str(e)}"
+                )
+                # Re-raise the exception after logging
+                raise
+
+    def get_all_aspects(
         self, from_createdon: datetime, stop_time: datetime
+    ) -> Iterable[Tuple[MetadataChangeProposalWrapper, datetime]]:
+        logger.info("Fetching Structured properties aspects")
+        yield from self.get_aspects(
+            from_createdon=from_createdon,
+            stop_time=stop_time,
+            set_structured_properties_filter=True,
+        )
+
+        logger.info(
+            f"Waiting for {self.config.structured_properties_template_cache_invalidation_interval} seconds for structured properties cache to invalidate"
+        )
+
+        time.sleep(
+            self.config.structured_properties_template_cache_invalidation_interval
+        )
+
+        logger.info("Fetching aspects")
+        yield from self.get_aspects(
+            from_createdon=from_createdon,
+            stop_time=stop_time,
+            set_structured_properties_filter=False,
+        )
+
+    def get_aspects(
+        self,
+        from_createdon: datetime,
+        stop_time: datetime,
+        set_structured_properties_filter: bool = False,
     ) -> Iterable[Tuple[MetadataChangeProposalWrapper, datetime]]:
         orderer = VersionOrderer[Dict[str, Any]](
             enabled=self.config.include_all_versions
         )
-        rows = self._get_rows(from_createdon=from_createdon, stop_time=stop_time)
+        rows = self._get_rows(
+            start_date=from_createdon,
+            end_date=stop_time,
+            set_structured_properties_filter=set_structured_properties_filter,
+            limit=self.config.database_query_batch_size,
+        )
         for row in orderer(rows):
             mcp = self._parse_row(row)
             if mcp:
@@ -190,22 +348,29 @@ class DataHubDatabaseReader:
 
     def get_soft_deleted_rows(self) -> Iterable[Dict[str, Any]]:
         """
-        Fetches all soft-deleted entities from the database.
+        Fetches all soft-deleted entities from the database using pagination.
 
         Yields:
             Row objects containing URNs of soft-deleted entities
         """
-        with self.engine.connect() as conn:
-            with contextlib.closing(conn.connection.cursor()) as cursor:
-                logger.debug("Polling soft-deleted urns from database")
-                cursor.execute(self.soft_deleted_urns_query)
-                columns = [desc[0] for desc in cursor.description]
-                while True:
-                    rows = cursor.fetchmany(self.config.database_query_batch_size)
-                    if not rows:
-                        return
-                    for row in rows:
-                        yield dict(zip(columns, row))
+        try:
+            params: Dict = {}
+
+            logger.debug("Fetching soft-deleted URNs")
+
+            # Use server-side cursor implementation
+            rows = self.execute_server_cursor(self.soft_deleted_urns_query, params)
+            processed_rows = 0
+            # Process and yield rows
+            for row in rows:
+                processed_rows += 1
+                yield row
+
+            logger.debug(f"Fetched batch of {processed_rows} soft-deleted URNs")
+
+        except Exception:
+            logger.exception("Error fetching soft-deleted row", exc_info=True)
+            raise
 
     def _parse_row(
         self, row: Dict[str, Any]
@@ -215,12 +380,16 @@ class DataHubDatabaseReader:
             json_metadata = post_json_transform(
                 json.loads(row["systemmetadata"] or "{}")
             )
-            system_metadata = SystemMetadataClass.from_obj(json_metadata)
+            system_metadata = None
+            if self.config.preserve_system_metadata:
+                system_metadata = SystemMetadataClass.from_obj(json_metadata)
+                if system_metadata.properties:
+                    is_no_op = system_metadata.properties.pop("isNoOp", None)
+                    logger.debug(f"Removed potential value for is_no_op={is_no_op}")
             return MetadataChangeProposalWrapper(
                 entityUrn=row["urn"],
                 aspect=ASPECT_MAP[row["aspect"]].from_obj(json_aspect),
                 systemMetadata=system_metadata,
-                changeType=ChangeTypeClass.UPSERT,
             )
         except Exception as e:
             logger.warning(

datahub/ingestion/source/datahub/datahub_source.py

@@ -6,7 +6,9 @@ from typing import Dict, Iterable, List, Optional
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
+    SourceCapability,
     SupportStatus,
+    capability,
     config_class,
     platform_name,
     support_status,
@@ -17,6 +19,7 @@ from datahub.ingestion.api.source_helpers import (
     auto_workunit_reporter,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.datahub.config import DataHubSourceConfig
 from datahub.ingestion.source.datahub.datahub_api_reader import DataHubApiReader
 from datahub.ingestion.source.datahub.datahub_database_reader import (
@@ -37,6 +40,13 @@ logger = logging.getLogger(__name__)
 @platform_name("DataHub")
 @config_class(DataHubSourceConfig)
 @support_status(SupportStatus.TESTING)
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+    ],
+)
 class DataHubSource(StatefulIngestionSourceBase):
     platform: str = "datahub"
 
@@ -117,7 +127,7 @@ class DataHubSource(StatefulIngestionSourceBase):
     ) -> Iterable[MetadataWorkUnit]:
         logger.info(f"Fetching database aspects starting from {from_createdon}")
         progress = ProgressTimer(report_every=timedelta(seconds=60))
-        mcps = reader.get_aspects(from_createdon, self.report.stop_time)
+        mcps = reader.get_all_aspects(from_createdon, self.report.stop_time)
         for i, (mcp, createdon) in enumerate(mcps):
             if not self.urn_pattern.allowed(str(mcp.entityUrn)):
                 continue

datahub/ingestion/source/dbt/dbt_cloud.py

@@ -9,6 +9,7 @@ import requests
 from pydantic import Field, root_validator
 
 from datahub.ingestion.api.decorators import (
+    SourceCapability,
     SupportStatus,
     capability,
     config_class,
@@ -17,7 +18,6 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.source import (
     CapabilityReport,
-    SourceCapability,
     TestableSource,
     TestConnectionReport,
 )
@@ -26,6 +26,7 @@ from datahub.ingestion.source.dbt.dbt_common import (
     DBTCommonConfig,
     DBTNode,
     DBTSourceBase,
+    DBTSourceReport,
 )
 from datahub.ingestion.source.dbt.dbt_tests import DBTTest, DBTTestResult
 
@@ -262,16 +263,16 @@ query DatahubMetadataQuery_{type}($jobId: BigInt!, $runId: BigInt) {{
 
 @platform_name("dbt")
 @config_class(DBTCloudConfig)
-@support_status(SupportStatus.INCUBATING)
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
-@capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
+@support_status(SupportStatus.CERTIFIED)
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class DBTCloudSource(DBTSourceBase, TestableSource):
     config: DBTCloudConfig
+    report: DBTSourceReport  # nothing cloud-specific in the report
 
     @classmethod
     def create(cls, config_dict, ctx):
         config = DBTCloudConfig.parse_obj(config_dict)
-        return cls(config, ctx, "dbt")
+        return cls(config, ctx)
 
     @staticmethod
     def test_connection(config_dict: dict) -> TestConnectionReport:
@@ -369,9 +370,12 @@ class DBTCloudSource(DBTSourceBase, TestableSource):
             name = node["alias"]
 
         comment = node.get("comment", "")
-        description = node["description"]
-        if node.get("sourceDescription"):
-            description = node["sourceDescription"]
+
+        # In dbt sources, there are two types of descriptions:
+        # - description: table-level description (specific to the source table)
+        # - sourceDescription: schema-level description (describes the overall source schema)
+        # The table-level description should take precedence since it's more specific.
+        description = node["description"] or node.get("sourceDescription", "")
 
         if node["resourceType"] == "model":
             materialization = node["materializedType"]
@@ -405,8 +409,11 @@ class DBTCloudSource(DBTSourceBase, TestableSource):
         if node["resourceType"] in {"model", "seed", "snapshot"}:
             status = node["status"]
             if status is None and materialization != "ephemeral":
-                self.report.report_warning(
-                    key, "node is missing a status, schema metadata will be incomplete"
+                self.report.warning(
+                    title="Schema information may be incomplete",
+                    message="Some nodes are missing the `status` field, which dbt uses to track the status of the node in the target database.",
+                    context=key,
+                    log=False,
                 )
 
             # The code fields are new in dbt 1.3, and replace the sql ones.