acryl-datahub 1.1.0.4rc2__py3-none-any.whl → 1.1.0.5__py3-none-any.whl

datahub/ingestion/source/dremio/dremio_sql_queries.py

@@ -1,3 +1,7 @@
+from datetime import datetime, timedelta
+from typing import Optional
+
+
 class DremioSQLQueries:
     QUERY_DATASETS_CE = """
     SELECT* FROM
@@ -235,28 +239,83 @@ class DremioSQLQueries:
             TABLE_NAME ASC
             """
 
-    # Dremio Documentation: https://docs.dremio.com/current/reference/sql/system-tables/jobs_recent/
-    # queried_datasets incorrectly documented as [varchar]. Observed as varchar.
-    # LENGTH used as opposed to ARRAY_SIZE
-    QUERY_ALL_JOBS = """
-    SELECT
-        job_id,
-        user_name,
-        submitted_ts,
-        query,
-        queried_datasets
-    FROM
-        SYS.JOBS_RECENT
-    WHERE
-        STATUS = 'COMPLETED'
-        AND LENGTH(queried_datasets)>0
-        AND user_name != '$dremio$'
-        AND query_type not like '%INTERNAL%'
-    """
+    @staticmethod
+    def _get_default_start_timestamp_millis() -> str:
+        """Get default start timestamp (1 day ago) in milliseconds precision format"""
+        one_day_ago = datetime.now() - timedelta(days=1)
+        return one_day_ago.strftime("%Y-%m-%d %H:%M:%S.%f")[
+            :-3
+        ]  # Truncate to milliseconds
+
+    @staticmethod
+    def _get_default_end_timestamp_millis() -> str:
+        """Get default end timestamp (now) in milliseconds precision format"""
+        now = datetime.now()
+        return now.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]  # Truncate to milliseconds
+
+    @staticmethod
+    def get_query_all_jobs(
+        start_timestamp_millis: Optional[str] = None,
+        end_timestamp_millis: Optional[str] = None,
+    ) -> str:
+        """
+        Get query for all jobs with optional time filtering.
+
+        Args:
+            start_timestamp_millis: Start timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to 1 day ago)
+            end_timestamp_millis: End timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to now)
+
+        Returns:
+            SQL query string with time filtering applied
+        """
+        if start_timestamp_millis is None:
+            start_timestamp_millis = (
+                DremioSQLQueries._get_default_start_timestamp_millis()
+            )
+        if end_timestamp_millis is None:
+            end_timestamp_millis = DremioSQLQueries._get_default_end_timestamp_millis()
+
+        return f"""
+        SELECT
+            job_id,
+            user_name,
+            submitted_ts,
+            query,
+            queried_datasets
+        FROM
+            SYS.JOBS_RECENT
+        WHERE
+            STATUS = 'COMPLETED'
+            AND LENGTH(queried_datasets)>0
+            AND user_name != '$dremio$'
+            AND query_type not like '%INTERNAL%'
+            AND submitted_ts >= TIMESTAMP '{start_timestamp_millis}'
+            AND submitted_ts <= TIMESTAMP '{end_timestamp_millis}'
+        """
+
+    @staticmethod
+    def get_query_all_jobs_cloud(
+        start_timestamp_millis: Optional[str] = None,
+        end_timestamp_millis: Optional[str] = None,
+    ) -> str:
+        """
+        Get query for all jobs in Dremio Cloud with optional time filtering.
+
+        Args:
+            start_timestamp_millis: Start timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to 7 days ago)
+            end_timestamp_millis: End timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to now)
+
+        Returns:
+            SQL query string with time filtering applied
+        """
+        if start_timestamp_millis is None:
+            start_timestamp_millis = (
+                DremioSQLQueries._get_default_start_timestamp_millis()
+            )
+        if end_timestamp_millis is None:
+            end_timestamp_millis = DremioSQLQueries._get_default_end_timestamp_millis()
 
-    # Dremio Documentation: https://docs.dremio.com/cloud/reference/sql/system-tables/jobs-historical
-    # queried_datasets correctly documented as [varchar]
-    QUERY_ALL_JOBS_CLOUD = """
+        return f"""
         SELECT
             job_id,
             user_name,
@@ -270,6 +329,8 @@ class DremioSQLQueries:
             AND ARRAY_SIZE(queried_datasets)>0
             AND user_name != '$dremio$'
             AND query_type not like '%INTERNAL%'
+            AND submitted_ts >= TIMESTAMP '{start_timestamp_millis}'
+            AND submitted_ts <= TIMESTAMP '{end_timestamp_millis}'
         """
 
     QUERY_TYPES = [

datahub/ingestion/source/fivetran/fivetran.py

@@ -1,8 +1,8 @@
 import logging
-from typing import Dict, Iterable, List, Optional
+from typing import Dict, Iterable, List, Optional, Union
 
 import datahub.emitter.mce_builder as builder
-from datahub.api.entities.datajob import DataFlow, DataJob
+from datahub.api.entities.datajob import DataJob as DataJobV1
 from datahub.api.entities.dataprocess.dataprocess_instance import (
     DataProcessInstance,
     InstanceRunResult,
@@ -42,8 +42,10 @@ from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     FineGrainedLineageDownstreamType,
     FineGrainedLineageUpstreamType,
 )
-from datahub.utilities.urns.data_flow_urn import DataFlowUrn
-from datahub.utilities.urns.dataset_urn import DatasetUrn
+from datahub.metadata.urns import CorpUserUrn, DataFlowUrn, DatasetUrn
+from datahub.sdk.dataflow import DataFlow
+from datahub.sdk.datajob import DataJob
+from datahub.sdk.entity import Entity
 
 # Logger instance
 logger = logging.getLogger(__name__)
@@ -75,8 +77,8 @@ class FivetranSource(StatefulIngestionSourceBase):
         self.audit_log = FivetranLogAPI(self.config.fivetran_log_config)
 
     def _extend_lineage(self, connector: Connector, datajob: DataJob) -> Dict[str, str]:
-        input_dataset_urn_list: List[DatasetUrn] = []
-        output_dataset_urn_list: List[DatasetUrn] = []
+        input_dataset_urn_list: List[Union[str, DatasetUrn]] = []
+        output_dataset_urn_list: List[Union[str, DatasetUrn]] = []
         fine_grained_lineage: List[FineGrainedLineage] = []
 
         # TODO: Once Fivetran exposes the database via the API, we shouldn't ask for it via config.
@@ -178,9 +180,9 @@ class FivetranSource(StatefulIngestionSourceBase):
                         )
                     )
 
-        datajob.inlets.extend(input_dataset_urn_list)
-        datajob.outlets.extend(output_dataset_urn_list)
-        datajob.fine_grained_lineages.extend(fine_grained_lineage)
+        datajob.set_inlets(input_dataset_urn_list)
+        datajob.set_outlets(output_dataset_urn_list)
+        datajob.set_fine_grained_lineages(fine_grained_lineage)
 
         return dict(
             **{
@@ -197,10 +199,10 @@ class FivetranSource(StatefulIngestionSourceBase):
 
     def _generate_dataflow_from_connector(self, connector: Connector) -> DataFlow:
         return DataFlow(
-            orchestrator=Constant.ORCHESTRATOR,
-            id=connector.connector_id,
+            platform=Constant.ORCHESTRATOR,
+            name=connector.connector_id,
             env=self.config.env,
-            name=connector.connector_name,
+            display_name=connector.connector_name,
             platform_instance=self.config.platform_instance,
         )
 
@@ -213,11 +215,11 @@ class FivetranSource(StatefulIngestionSourceBase):
         )
         owner_email = self.audit_log.get_user_email(connector.user_id)
         datajob = DataJob(
-            id=connector.connector_id,
+            name=connector.connector_id,
             flow_urn=dataflow_urn,
             platform_instance=self.config.platform_instance,
-            name=connector.connector_name,
-            owners={owner_email} if owner_email else set(),
+            display_name=connector.connector_name,
+            owners=[CorpUserUrn(owner_email)] if owner_email else None,
         )
 
         # Map connector source and destination table with dataset entity
@@ -232,16 +234,24 @@ class FivetranSource(StatefulIngestionSourceBase):
             "sync_frequency": str(connector.sync_frequency),
             "destination_id": connector.destination_id,
         }
-        datajob.properties = {
-            **connector_properties,
-            **lineage_properties,
-        }
+
+        datajob.set_custom_properties({**connector_properties, **lineage_properties})
 
         return datajob
 
     def _generate_dpi_from_job(self, job: Job, datajob: DataJob) -> DataProcessInstance:
+        # hack: convert to old instance for DataProcessInstance.from_datajob compatibility
+        datajob_v1 = DataJobV1(
+            id=datajob.name,
+            flow_urn=datajob.flow_urn,
+            platform_instance=self.config.platform_instance,
+            name=datajob.name,
+            inlets=datajob.inlets,
+            outlets=datajob.outlets,
+            fine_grained_lineages=datajob.fine_grained_lineages,
+        )
         return DataProcessInstance.from_datajob(
-            datajob=datajob,
+            datajob=datajob_v1,
             id=job.job_id,
             clone_inlets=True,
             clone_outlets=True,
@@ -278,17 +288,15 @@ class FivetranSource(StatefulIngestionSourceBase):
 
     def _get_connector_workunits(
         self, connector: Connector
-    ) -> Iterable[MetadataWorkUnit]:
+    ) -> Iterable[Union[MetadataWorkUnit, Entity]]:
         self.report.report_connectors_scanned()
         # Create dataflow entity with same name as connector name
         dataflow = self._generate_dataflow_from_connector(connector)
-        for mcp in dataflow.generate_mcp():
-            yield mcp.as_workunit()
+        yield dataflow
 
         # Map Fivetran's connector entity with Datahub's datajob entity
         datajob = self._generate_datajob_from_connector(connector)
-        for mcp in datajob.generate_mcp(materialize_iolets=False):
-            yield mcp.as_workunit()
+        yield datajob
 
         # Map Fivetran's job/sync history entity with Datahub's data process entity
         if len(connector.jobs) >= MAX_JOBS_PER_CONNECTOR:
@@ -310,7 +318,7 @@ class FivetranSource(StatefulIngestionSourceBase):
             ).workunit_processor,
         ]
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
         """
         Datahub Ingestion framework invoke this method
         """

datahub/ingestion/source/gcs/gcs_source.py

@@ -16,6 +16,7 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceCapability
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.aws_common import AwsConnectionConfig
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.data_lake_common.config import PathSpecsConfigMixin
 from datahub.ingestion.source.data_lake_common.data_lake_utils import PLATFORM_GCS
 from datahub.ingestion.source.data_lake_common.object_store import (
@@ -82,7 +83,14 @@ class GCSSourceReport(DataLakeSourceReport):
 @platform_name("Google Cloud Storage", id=PLATFORM_GCS)
 @config_class(GCSSourceConfig)
 @support_status(SupportStatus.INCUBATING)
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.GCS_BUCKET,
+        SourceCapabilityModifier.FOLDER,
+    ],
+)
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
 @capability(SourceCapability.DATA_PROFILING, "Not supported", supported=False)
 class GCSSource(StatefulIngestionSourceBase):
@@ -112,6 +120,7 @@ class GCSSource(StatefulIngestionSourceBase):
             env=self.config.env,
             max_rows=self.config.max_rows,
             number_of_files_to_sample=self.config.number_of_files_to_sample,
+            platform=PLATFORM_GCS,  # Ensure GCS platform is used for correct container subtypes
         )
         return s3_config
 
@@ -138,7 +147,9 @@ class GCSSource(StatefulIngestionSourceBase):
 
     def create_equivalent_s3_source(self, ctx: PipelineContext) -> S3Source:
         config = self.create_equivalent_s3_config()
-        s3_source = S3Source(config, PipelineContext(ctx.run_id))
+        # Create a new context for S3 source without graph to avoid duplicate checkpointer registration
+        s3_ctx = PipelineContext(run_id=ctx.run_id, pipeline_name=ctx.pipeline_name)
+        s3_source = S3Source(config, s3_ctx)
         return self.s3_source_overrides(s3_source)
 
     def s3_source_overrides(self, source: S3Source) -> S3Source:

datahub/ingestion/source/ge_data_profiler.py

@@ -120,7 +120,6 @@ SNOWFLAKE = "snowflake"
 BIGQUERY = "bigquery"
 REDSHIFT = "redshift"
 DATABRICKS = "databricks"
-TRINO = "trino"
 
 # Type names for Databricks, to match Title Case types in sqlalchemy
 ProfilerTypeMapping.INT_TYPE_NAMES.append("Integer")
@@ -206,6 +205,17 @@ def get_column_unique_count_dh_patch(self: SqlAlchemyDataset, column: str) -> in
             )
         )
         return convert_to_json_serializable(element_values.fetchone()[0])
+    elif (
+        self.engine.dialect.name.lower() == GXSqlDialect.AWSATHENA
+        or self.engine.dialect.name.lower() == GXSqlDialect.TRINO
+    ):
+        return convert_to_json_serializable(
+            self.engine.execute(
+                sa.select(sa.func.approx_distinct(sa.column(column))).select_from(
+                    self._table
+                )
+            ).scalar()
+        )
     return convert_to_json_serializable(
         self.engine.execute(
             sa.select([sa.func.count(sa.func.distinct(sa.column(column)))]).select_from(
@@ -734,11 +744,41 @@ class _SingleDatasetProfiler(BasicDatasetProfilerBase):
     def _get_dataset_column_distinct_value_frequencies(
         self, column_profile: DatasetFieldProfileClass, column: str
     ) -> None:
-        if self.config.include_field_distinct_value_frequencies:
+        if not self.config.include_field_distinct_value_frequencies:
+            return
+        try:
+            results = self.dataset.engine.execute(
+                sa.select(
+                    [
+                        sa.column(column),
+                        sa.func.count(sa.column(column)),
+                    ]
+                )
+                .select_from(self.dataset._table)
+                .where(sa.column(column).is_not(None))
+                .group_by(sa.column(column))
+            ).fetchall()
+
             column_profile.distinctValueFrequencies = [
-                ValueFrequencyClass(value=str(value), frequency=count)
-                for value, count in self.dataset.get_column_value_counts(column).items()
+                ValueFrequencyClass(value=str(value), frequency=int(count))
+                for value, count in results
             ]
+            # sort so output is deterministic. don't do it in SQL because not all column
+            # types are sortable in SQL (such as JSON data types on Athena/Trino).
+            column_profile.distinctValueFrequencies = sorted(
+                column_profile.distinctValueFrequencies, key=lambda x: x.value
+            )
+        except Exception as e:
+            logger.debug(
+                f"Caught exception while attempting to get distinct value frequencies for column {column}. {e}"
+            )
+
+            self.report.report_warning(
+                title="Profiling: Unable to Calculate Distinct Value Frequencies",
+                message="Distinct value frequencies for the column will not be accessible",
+                context=f"{self.dataset_name}.{column}",
+                exc=e,
+            )
 
     @_run_with_query_combiner
     def _get_dataset_column_histogram(
@@ -1173,26 +1213,34 @@ class DatahubGEProfiler:
             f"Will profile {len(requests)} table(s) with {max_workers} worker(s) - this may take a while"
         )
 
-        with PerfTimer() as timer, unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_unique_count",
-            get_column_unique_count_dh_patch,
-        ), unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_bigquery",
-            _get_column_quantiles_bigquery_patch,
-        ), unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_awsathena",
-            _get_column_quantiles_awsathena_patch,
-        ), unittest.mock.patch(
-            "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_median",
-            _get_column_median_patch,
-        ), concurrent.futures.ThreadPoolExecutor(
-            max_workers=max_workers
-        ) as async_executor, SQLAlchemyQueryCombiner(
-            enabled=self.config.query_combiner_enabled,
-            catch_exceptions=self.config.catch_exceptions,
-            is_single_row_query_method=_is_single_row_query_method,
-            serial_execution_fallback_enabled=True,
-        ).activate() as query_combiner:
+        with (
+            PerfTimer() as timer,
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_unique_count",
+                get_column_unique_count_dh_patch,
+            ),
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_bigquery",
+                _get_column_quantiles_bigquery_patch,
+            ),
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset._get_column_quantiles_awsathena",
+                _get_column_quantiles_awsathena_patch,
+            ),
+            unittest.mock.patch(
+                "great_expectations.dataset.sqlalchemy_dataset.SqlAlchemyDataset.get_column_median",
+                _get_column_median_patch,
+            ),
+            concurrent.futures.ThreadPoolExecutor(
+                max_workers=max_workers
+            ) as async_executor,
+            SQLAlchemyQueryCombiner(
+                enabled=self.config.query_combiner_enabled,
+                catch_exceptions=self.config.catch_exceptions,
+                is_single_row_query_method=_is_single_row_query_method,
+                serial_execution_fallback_enabled=True,
+            ).activate() as query_combiner,
+        ):
             # Submit the profiling requests to the thread pool executor.
             async_profiles = collections.deque(
                 async_executor.submit(
@@ -1395,12 +1443,12 @@ class DatahubGEProfiler:
                     )
                 return None
             finally:
-                if batch is not None and self.base_engine.engine.name.upper() in [
-                    "TRINO",
-                    "AWSATHENA",
+                if batch is not None and self.base_engine.engine.name.lower() in [
+                    GXSqlDialect.TRINO,
+                    GXSqlDialect.AWSATHENA,
                 ]:
                     if (
-                        self.base_engine.engine.name.upper() == "TRINO"
+                        self.base_engine.engine.name.lower() == GXSqlDialect.TRINO
                         or temp_view is not None
                     ):
                         self._drop_temp_table(batch)

datahub/ingestion/source/hex/api.py

@@ -5,7 +5,9 @@ from typing import Any, Dict, Generator, List, Optional, Union
 
 import requests
 from pydantic import BaseModel, Field, ValidationError, validator
+from requests.adapters import HTTPAdapter
 from typing_extensions import assert_never
+from urllib3.util.retry import Retry
 
 from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.source.hex.constants import (
@@ -220,6 +222,7 @@ class HexApi:
         self.base_url = base_url
         self.report = report
         self.page_size = page_size
+        self.session = self._create_retry_session()
 
     def _list_projects_url(self):
         return f"{self.base_url}/projects"
@@ -227,6 +230,28 @@ class HexApi:
     def _auth_header(self):
         return {"Authorization": f"Bearer {self.token}"}
 
+    def _create_retry_session(self) -> requests.Session:
+        """Create a requests session with retry logic for rate limiting.
+
+        Hex API rate limit: 60 requests per minute
+        https://learn.hex.tech/docs/api/api-overview#kernel-and-rate-limits
+        """
+        session = requests.Session()
+
+        # Configure retry strategy for 429 (Too Many Requests) with exponential backoff
+        retry_strategy = Retry(
+            total=5,  # Maximum number of retries
+            status_forcelist=[429],  # Only retry on 429 status code
+            backoff_factor=2,  # Exponential backoff: 2, 4, 8, 16, 32 seconds
+            raise_on_status=True,  # Raise exception after max retries
+        )
+
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        session.mount("http://", adapter)
+        session.mount("https://", adapter)
+
+        return session
+
     def fetch_projects(
         self,
         include_components: bool = True,
@@ -259,7 +284,7 @@ class HexApi:
         logger.debug(f"Fetching projects page with params: {params}")
         self.report.fetch_projects_page_calls += 1
         try:
-            response = requests.get(
+            response = self.session.get(
                 url=self._list_projects_url(),
                 headers=self._auth_header(),
                 params=params,

datahub/ingestion/source/identity/azure_ad.py

@@ -167,7 +167,7 @@ class AzureADSourceReport(StaleEntityRemovalSourceReport):
 @config_class(AzureADConfig)
 @support_status(SupportStatus.CERTIFIED)
 @capability(
-    SourceCapability.DELETION_DETECTION, "Optionally enabled via stateful_ingestion"
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 class AzureADSource(StatefulIngestionSourceBase):
     """

datahub/ingestion/source/identity/okta.py

@@ -41,7 +41,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import (
 )
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
-    ChangeTypeClass,
     CorpGroupInfoClass,
     CorpUserInfoClass,
     GroupMembershipClass,
@@ -202,7 +201,7 @@ class OktaSourceReport(StaleEntityRemovalSourceReport):
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.DESCRIPTIONS, "Optionally enabled via configuration")
 @capability(
-    SourceCapability.DELETION_DETECTION, "Optionally enabled via stateful_ingestion"
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 class OktaSource(StatefulIngestionSourceBase):
     """
@@ -332,18 +331,12 @@ class OktaSource(StatefulIngestionSourceBase):
                 yield MetadataWorkUnit(id=wu_id, mce=mce)
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpGroup",
                     entityUrn=datahub_corp_group_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="origin",
                     aspect=OriginClass(OriginTypeClass.EXTERNAL, "OKTA"),
                 ).as_workunit()
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpGroup",
                     entityUrn=datahub_corp_group_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="status",
                     aspect=StatusClass(removed=False),
                 ).as_workunit()
 
@@ -418,18 +411,12 @@ class OktaSource(StatefulIngestionSourceBase):
                 yield MetadataWorkUnit(id=wu_id, mce=mce)
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpuser",
                     entityUrn=datahub_corp_user_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="origin",
                     aspect=OriginClass(OriginTypeClass.EXTERNAL, "OKTA"),
                 ).as_workunit()
 
                 yield MetadataChangeProposalWrapper(
-                    entityType="corpuser",
                     entityUrn=datahub_corp_user_snapshot.urn,
-                    changeType=ChangeTypeClass.UPSERT,
-                    aspectName="status",
                     aspect=StatusClass(removed=False),
                 ).as_workunit()