acryl-datahub 1.3.0.1rc5__py3-none-any.whl → 1.3.0.1rc7__py3-none-any.whl

datahub/ingestion/source/metabase.py

@@ -52,6 +52,7 @@ from datahub.metadata.schema_classes import (
     ChartQueryTypeClass,
     ChartTypeClass,
     DashboardInfoClass,
+    EdgeClass,
     OwnerClass,
     OwnershipClass,
     OwnershipTypeClass,
@@ -338,19 +339,25 @@ class MetabaseSource(StatefulIngestionSourceBase):
             lastModified=AuditStamp(time=modified_ts, actor=modified_actor),
         )
 
-        chart_urns = []
+        # Convert chart URNs to chart edges (instead of deprecated charts field)
+        chart_edges = []
         cards_data = dashboard_details.get("dashcards", {})
         for card_info in cards_data:
             card_id = card_info.get("card").get("id", "")
             if not card_id:
                 continue  # most likely a virtual card without an id (text or heading), not relevant.
             chart_urn = builder.make_chart_urn(self.platform, str(card_id))
-            chart_urns.append(chart_urn)
+            chart_edges.append(
+                EdgeClass(
+                    destinationUrn=chart_urn,
+                    lastModified=last_modified.lastModified,
+                )
+            )
 
         dashboard_info_class = DashboardInfoClass(
             description=description,
             title=title,
-            charts=chart_urns,
+            chartEdges=chart_edges,
             lastModified=last_modified,
             dashboardUrl=f"{self.config.display_uri}/dashboard/{dashboard_id}",
             customProperties={},
@@ -488,13 +495,25 @@ class MetabaseSource(StatefulIngestionSourceBase):
         datasource_urn = self.get_datasource_urn(card_details)
         custom_properties = self.construct_card_custom_properties(card_details)
 
+        input_edges = (
+            [
+                EdgeClass(
+                    destinationUrn=urn,
+                    lastModified=last_modified.lastModified,
+                )
+                for urn in datasource_urn
+            ]
+            if datasource_urn
+            else None
+        )
+
         chart_info = ChartInfoClass(
             type=chart_type,
             description=description,
             title=title,
             lastModified=last_modified,
             chartUrl=f"{self.config.display_uri}/card/{card_id}",
-            inputs=datasource_urn,
+            inputEdges=input_edges,
             customProperties=custom_properties,
         )
         chart_snapshot.aspects.append(chart_info)

datahub/ingestion/source/mlflow.py

@@ -136,7 +136,7 @@ class MLflowRegisteredModelStageInfo:
 
 @platform_name("MLflow")
 @config_class(MLflowConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @capability(
     SourceCapability.DESCRIPTIONS,
     "Extract descriptions for MLflow Registered Models and Model Versions",

datahub/ingestion/source/s3/source.py

@@ -188,7 +188,7 @@ class TableData:
 
 @platform_name("S3 / Local Files", id="s3")
 @config_class(DataLakeSourceConfig)
-@support_status(SupportStatus.INCUBATING)
+@support_status(SupportStatus.CERTIFIED)
 @capability(
     SourceCapability.CONTAINERS,
     "Enabled by default",

datahub/ingestion/source/salesforce.py

@@ -527,7 +527,7 @@ class SalesforceApi:
 
 @platform_name("Salesforce")
 @config_class(SalesforceConfig)
-@support_status(SupportStatus.INCUBATING)
+@support_status(SupportStatus.CERTIFIED)
 @capability(
     capability_name=SourceCapability.PLATFORM_INSTANCE,
     description="Can be equivalent to Salesforce organization",

datahub/ingestion/source/slack/slack.py

@@ -245,7 +245,7 @@ DATA_PLATFORM_SLACK_URN: str = builder.make_data_platform_urn(PLATFORM_NAME)
 
 @platform_name("Slack")
 @config_class(SlackSourceConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.CERTIFIED)
 class SlackSource(StatefulIngestionSourceBase):
     def __init__(self, ctx: PipelineContext, config: SlackSourceConfig):
         super().__init__(config, ctx)

datahub/ingestion/source/snowflake/snowflake_config.py

@@ -31,6 +31,7 @@ from datahub.ingestion.source.sql.sql_config import SQLCommonConfig, SQLFilterCo
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulLineageConfigMixin,
     StatefulProfilingConfigMixin,
+    StatefulTimeWindowConfigMixin,
     StatefulUsageConfigMixin,
 )
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
@@ -199,6 +200,7 @@ class SnowflakeV2Config(
     SnowflakeUsageConfig,
     StatefulLineageConfigMixin,
     StatefulUsageConfigMixin,
+    StatefulTimeWindowConfigMixin,
     StatefulProfilingConfigMixin,
     ClassificationSourceConfigMixin,
     IncrementalPropertiesConfigMixin,
@@ -477,6 +479,20 @@ class SnowflakeV2Config(
 
         return shares
 
+    @root_validator(pre=False, skip_on_failure=True)
+    def validate_queries_v2_stateful_ingestion(cls, values: Dict) -> Dict:
+        if values.get("use_queries_v2"):
+            if values.get("enable_stateful_lineage_ingestion") or values.get(
+                "enable_stateful_usage_ingestion"
+            ):
+                logger.warning(
+                    "enable_stateful_lineage_ingestion and enable_stateful_usage_ingestion are deprecated "
+                    "when using use_queries_v2=True. These configs only work with the legacy (non-queries v2) extraction path. "
+                    "For queries v2, use enable_stateful_time_window instead to enable stateful ingestion "
+                    "for the unified time window extraction (lineage + usage + operations + queries)."
+                )
+        return values
+
     def outbounds(self) -> Dict[str, Set[DatabaseId]]:
         """
         Returns mapping of

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -17,9 +17,11 @@ from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFr
 from datahub.configuration.time_window_config import (
     BaseTimeWindowConfig,
     BucketDuration,
+    get_time_bucket,
 )
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import SupportStatus, config_class, support_status
 from datahub.ingestion.api.report import Report
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.source_helpers import auto_workunit
@@ -50,6 +52,9 @@ from datahub.ingestion.source.snowflake.stored_proc_lineage import (
     StoredProcLineageReport,
     StoredProcLineageTracker,
 )
+from datahub.ingestion.source.state.redundant_run_skip_handler import (
+    RedundantQueriesRunSkipHandler,
+)
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
 from datahub.metadata.urns import CorpUserUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
@@ -180,6 +185,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         structured_report: SourceReport,
         filters: SnowflakeFilter,
         identifiers: SnowflakeIdentifierBuilder,
+        redundant_run_skip_handler: Optional[RedundantQueriesRunSkipHandler] = None,
         graph: Optional[DataHubGraph] = None,
         schema_resolver: Optional[SchemaResolver] = None,
         discovered_tables: Optional[List[str]] = None,
@@ -191,9 +197,13 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         self.filters = filters
         self.identifiers = identifiers
         self.discovered_tables = set(discovered_tables) if discovered_tables else None
+        self.redundant_run_skip_handler = redundant_run_skip_handler
 
         self._structured_report = structured_report
 
+        # Adjust time window based on stateful ingestion state
+        self.start_time, self.end_time = self._get_time_window()
+
         # The exit stack helps ensure that we close all the resources we open.
         self._exit_stack = contextlib.ExitStack()
 
@@ -211,8 +221,8 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                 generate_query_usage_statistics=self.config.include_query_usage_statistics,
                 usage_config=BaseUsageConfig(
                     bucket_duration=self.config.window.bucket_duration,
-                    start_time=self.config.window.start_time,
-                    end_time=self.config.window.end_time,
+                    start_time=self.start_time,
+                    end_time=self.end_time,
                     user_email_pattern=self.config.user_email_pattern,
                     # TODO make the rest of the fields configurable
                 ),
@@ -228,6 +238,34 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
     def structured_reporter(self) -> SourceReport:
         return self._structured_report
 
+    def _get_time_window(self) -> tuple[datetime, datetime]:
+        if self.redundant_run_skip_handler:
+            start_time, end_time = (
+                self.redundant_run_skip_handler.suggest_run_time_window(
+                    self.config.window.start_time,
+                    self.config.window.end_time,
+                )
+            )
+        else:
+            start_time = self.config.window.start_time
+            end_time = self.config.window.end_time
+
+        # Usage statistics are aggregated per bucket (typically per day).
+        # To ensure accurate aggregated metrics, we need to align the start_time
+        # to the beginning of a bucket so that we include complete bucket periods.
+        if self.config.include_usage_statistics:
+            start_time = get_time_bucket(start_time, self.config.window.bucket_duration)
+
+        return start_time, end_time
+
+    def _update_state(self) -> None:
+        if self.redundant_run_skip_handler:
+            self.redundant_run_skip_handler.update_state(
+                self.config.window.start_time,
+                self.config.window.end_time,
+                self.config.window.bucket_duration,
+            )
+
     @functools.cached_property
     def local_temp_path(self) -> pathlib.Path:
         if self.config.local_temp_path:
@@ -355,6 +393,9 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         with self.report.aggregator_generate_timer:
             yield from auto_workunit(self.aggregator.gen_metadata())
 
+        # Update the stateful ingestion state after successful extraction
+        self._update_state()
+
     def fetch_users(self) -> UsersMapping:
         users: UsersMapping = dict()
         with self.structured_reporter.report_exc("Error fetching users from Snowflake"):
@@ -378,8 +419,8 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         # Derived from _populate_external_lineage_from_copy_history.
 
         query: str = SnowflakeQuery.copy_lineage_history(
-            start_time_millis=int(self.config.window.start_time.timestamp() * 1000),
-            end_time_millis=int(self.config.window.end_time.timestamp() * 1000),
+            start_time_millis=int(self.start_time.timestamp() * 1000),
+            end_time_millis=int(self.end_time.timestamp() * 1000),
             downstreams_deny_pattern=self.config.temporary_tables_pattern,
         )
 
@@ -414,8 +455,8 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         Union[PreparsedQuery, TableRename, TableSwap, ObservedQuery, StoredProcCall]
     ]:
         query_log_query = QueryLogQueryBuilder(
-            start_time=self.config.window.start_time,
-            end_time=self.config.window.end_time,
+            start_time=self.start_time,
+            end_time=self.end_time,
             bucket_duration=self.config.window.bucket_duration,
             deny_usernames=self.config.pushdown_deny_usernames,
             allow_usernames=self.config.pushdown_allow_usernames,
@@ -710,6 +751,8 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         self._exit_stack.close()
 
 
+@support_status(SupportStatus.CERTIFIED)
+@config_class(SnowflakeQueriesSourceConfig)
 class SnowflakeQueriesSource(Source):
     def __init__(self, ctx: PipelineContext, config: SnowflakeQueriesSourceConfig):
         self.ctx = ctx

datahub/ingestion/source/snowflake/snowflake_summary.py

@@ -59,7 +59,7 @@ class SnowflakeSummaryReport(SourceReport, BaseTimeWindowReport):
 
 
 @config_class(SnowflakeSummaryConfig)
-@support_status(SupportStatus.INCUBATING)
+@support_status(SupportStatus.CERTIFIED)
 class SnowflakeSummarySource(Source):
     def __init__(self, ctx: PipelineContext, config: SnowflakeSummaryConfig):
         super().__init__(ctx)

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -73,6 +73,7 @@ from datahub.ingestion.source.snowflake.snowflake_utils import (
 from datahub.ingestion.source.state.profiling_state_handler import ProfilingHandler
 from datahub.ingestion.source.state.redundant_run_skip_handler import (
     RedundantLineageRunSkipHandler,
+    RedundantQueriesRunSkipHandler,
     RedundantUsageRunSkipHandler,
 )
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
@@ -207,7 +208,7 @@ class SnowflakeV2Source(
         )
         self.report.sql_aggregator = self.aggregator.report
 
-        if self.config.include_table_lineage:
+        if self.config.include_table_lineage and not self.config.use_queries_v2:
             redundant_lineage_run_skip_handler: Optional[
                 RedundantLineageRunSkipHandler
             ] = None
@@ -589,6 +590,17 @@ class SnowflakeV2Source(
             with self.report.new_stage(f"*: {QUERIES_EXTRACTION}"):
                 schema_resolver = self.aggregator._schema_resolver
 
+                redundant_queries_run_skip_handler: Optional[
+                    RedundantQueriesRunSkipHandler
+                ] = None
+                if self.config.enable_stateful_time_window:
+                    redundant_queries_run_skip_handler = RedundantQueriesRunSkipHandler(
+                        source=self,
+                        config=self.config,
+                        pipeline_name=self.ctx.pipeline_name,
+                        run_id=self.ctx.run_id,
+                    )
+
                 queries_extractor = SnowflakeQueriesExtractor(
                     connection=self.connection,
                     # TODO: this should be its own section in main recipe
@@ -614,6 +626,7 @@ class SnowflakeV2Source(
                     structured_report=self.report,
                     filters=self.filters,
                     identifiers=self.identifiers,
+                    redundant_run_skip_handler=redundant_queries_run_skip_handler,
                     schema_resolver=schema_resolver,
                     discovered_tables=self.discovered_datasets,
                     graph=self.ctx.graph,

datahub/ingestion/source/sql_queries.py

@@ -93,7 +93,7 @@ class SqlQueriesSourceReport(SourceReport):
     sql_aggregator: Optional[SqlAggregatorReport] = None
 
 
-@platform_name("SQL Queries")
+@platform_name("SQL Queries", id="sql-queries")
 @config_class(SqlQueriesSourceConfig)
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.LINEAGE_COARSE, "Parsed from SQL queries")

datahub/ingestion/source/state/redundant_run_skip_handler.py

@@ -244,3 +244,24 @@ class RedundantUsageRunSkipHandler(RedundantRunSkipHandler):
             cur_state.begin_timestamp_millis = datetime_to_ts_millis(start_time)
             cur_state.end_timestamp_millis = datetime_to_ts_millis(end_time)
             cur_state.bucket_duration = bucket_duration
+
+
+class RedundantQueriesRunSkipHandler(RedundantRunSkipHandler):
+    """
+    Handler for stateful ingestion of queries v2 extraction.
+    Manages the time window for audit log extraction that combines
+    lineage, usage, operations, and queries.
+    """
+
+    def get_job_name_suffix(self):
+        return "_audit_window"
+
+    def update_state(
+        self, start_time: datetime, end_time: datetime, bucket_duration: BucketDuration
+    ) -> None:
+        cur_checkpoint = self.get_current_checkpoint()
+        if cur_checkpoint:
+            cur_state = cast(BaseTimeWindowCheckpointState, cur_checkpoint.state)
+            cur_state.begin_timestamp_millis = datetime_to_ts_millis(start_time)
+            cur_state.end_timestamp_millis = datetime_to_ts_millis(end_time)
+            cur_state.bucket_duration = bucket_duration

datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -101,7 +101,9 @@ class StatefulLineageConfigMixin(ConfigModel):
         default=True,
         description="Enable stateful lineage ingestion."
         " This will store lineage window timestamps after successful lineage ingestion. "
-        "and will not run lineage ingestion for same timestamps in subsequent run. ",
+        "and will not run lineage ingestion for same timestamps in subsequent run. "
+        "NOTE: This only works with use_queries_v2=False (legacy extraction path). "
+        "For queries v2, use enable_stateful_time_window instead.",
     )
 
     _store_last_lineage_extraction_timestamp = pydantic_renamed_field(
@@ -150,7 +152,9 @@ class StatefulUsageConfigMixin(BaseTimeWindowConfig):
         default=True,
         description="Enable stateful lineage ingestion."
         " This will store usage window timestamps after successful usage ingestion. "
-        "and will not run usage ingestion for same timestamps in subsequent run. ",
+        "and will not run usage ingestion for same timestamps in subsequent run. "
+        "NOTE: This only works with use_queries_v2=False (legacy extraction path). "
+        "For queries v2, use enable_stateful_time_window instead.",
     )
 
     _store_last_usage_extraction_timestamp = pydantic_renamed_field(
@@ -169,6 +173,30 @@ class StatefulUsageConfigMixin(BaseTimeWindowConfig):
         return values
 
 
+class StatefulTimeWindowConfigMixin(BaseTimeWindowConfig):
+    enable_stateful_time_window: bool = Field(
+        default=False,
+        description="Enable stateful time window tracking."
+        " This will store the time window after successful extraction "
+        "and adjust the time window in subsequent runs to avoid reprocessing. "
+        "NOTE: This is ONLY applicable when using queries v2 (use_queries_v2=True). "
+        "This replaces enable_stateful_lineage_ingestion and enable_stateful_usage_ingestion "
+        "for the queries v2 extraction path, since queries v2 extracts lineage, usage, operations, "
+        "and queries together from a single audit log and uses a unified time window.",
+    )
+
+    @root_validator(skip_on_failure=True)
+    def time_window_stateful_option_validator(cls, values: Dict) -> Dict:
+        sti = values.get("stateful_ingestion")
+        if not sti or not sti.enabled:
+            if values.get("enable_stateful_time_window"):
+                logger.warning(
+                    "Stateful ingestion is disabled, disabling enable_stateful_time_window config option as well"
+                )
+                values["enable_stateful_time_window"] = False
+        return values
+
+
 @dataclass
 class StatefulIngestionReport(SourceReport):
     pass

datahub/ingestion/source/unity/source.py

@@ -176,7 +176,7 @@ logger: logging.Logger = logging.getLogger(__name__)
     supported=True,
 )
 @capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
-@support_status(SupportStatus.INCUBATING)
+@support_status(SupportStatus.CERTIFIED)
 class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
     """
     This plugin extracts the following metadata from Databricks Unity Catalog:

datahub/ingestion/source/vertexai/vertexai.py

@@ -145,7 +145,7 @@ class PipelineMetadata:
 
 @platform_name("Vertex AI", id="vertexai")
 @config_class(VertexAIConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @capability(
     SourceCapability.DESCRIPTIONS,
     "Extract descriptions for Vertex AI Registered Models and Model Versions",