acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__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/dynamodb/dynamodb.py

@@ -12,7 +12,7 @@ from typing import (
     Union,
 )
 
-from pydantic.fields import Field
+from pydantic import Field, PositiveInt
 
 from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.source_common import DatasetSourceConfigMixin
@@ -73,7 +73,6 @@ from datahub.utilities.registries.domain_registry import DomainRegistry
 
 MAX_ITEMS_TO_RETRIEVE = 100
 PAGE_SIZE = 100
-MAX_SCHEMA_SIZE = 300
 MAX_PRIMARY_KEYS_SIZE = 100
 FIELD_DELIMITER = "."
 
@@ -107,6 +106,10 @@ class DynamoDBConfig(
         'Refer "Advanced Configurations" section for more details',
     )
 
+    max_schema_size: PositiveInt = Field(
+        default=300, description="Maximum number of fields to include in the schema."
+    )
+
     table_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns for tables to filter in ingestion. The table name format is 'region.table'",
@@ -160,7 +163,7 @@ _attribute_type_to_field_type_mapping: Dict[str, Type] = {
 
 @platform_name("DynamoDB", id="dynamodb")
 @config_class(DynamoDBConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @capability(
     SourceCapability.PLATFORM_INSTANCE,
     "By default, platform_instance will use the AWS account id",
@@ -362,7 +365,7 @@ class DynamoDBSource(StatefulIngestionSourceBase):
         if self.config.include_table_item is None:
             return
         dataset_name = f"{region}.{table_name}"
-        if dataset_name not in self.config.include_table_item.keys():
+        if dataset_name not in self.config.include_table_item:
             return
         primary_key_list = self.config.include_table_item.get(dataset_name)
         assert isinstance(primary_key_list, List)
@@ -455,51 +458,53 @@ class DynamoDBSource(StatefulIngestionSourceBase):
     ) -> SchemaMetadataClass:
         """ "
         To construct the schema metadata, it will first sort the schema by the occurrence of attribute names
-        in descending order and truncate the schema by MAX_SCHEMA_SIZE, and then start to construct the
+        in descending order and truncate the schema by max_schema_size, and then start to construct the
         schema metadata sorted by attribute name
         """
 
         canonical_schema: List[SchemaField] = []
         schema_size = len(schema.values())
         table_fields = list(schema.values())
-        if schema_size > MAX_SCHEMA_SIZE:
+        if schema_size > self.config.max_schema_size:
             # downsample the schema, using frequency as the sort key
             self.report.report_warning(
                 title="Schema Size Too Large",
-                message=f"Downsampling the table schema because MAX_SCHEMA_SIZE threshold is {MAX_SCHEMA_SIZE}",
+                message=f"Downsampling the table schema because `max_schema_size` threshold is {self.config.max_schema_size}",
                 context=f"Collection: {dataset_urn}",
             )
 
             # Add this information to the custom properties so user can know they are looking at down sampled schema
             dataset_properties.customProperties["schema.downsampled"] = "True"
             dataset_properties.customProperties["schema.totalFields"] = f"{schema_size}"
-        # append each schema field, schema will be sorted by count descending and delimited_name ascending and sliced to only include MAX_SCHEMA_SIZE items
+        # append each schema field, schema will be sorted by count descending and delimited_name ascending and sliced to only include max_schema_size items
+        primary_keys = []
         for schema_field in sorted(
             table_fields,
             key=lambda x: (
                 -x["count"],
                 x["delimited_name"],
             ),  # Negate `count` for descending order, `delimited_name` stays the same for ascending
-        )[0:MAX_SCHEMA_SIZE]:
+        )[: self.config.max_schema_size]:
             field_path = schema_field["delimited_name"]
             native_data_type = self.get_native_type(schema_field["type"], table_name)
             type = self.get_field_type(schema_field["type"], table_name)
-            description = None
             nullable = True
             if field_path in primary_key_dict:
-                description = (
+                # primary key should not be nullable
+                type_key = (
                     "Partition Key"
                     if primary_key_dict.get(field_path) == "HASH"
                     else "Sort Key"
                 )
-                # primary key should not be nullable
+                dataset_properties.customProperties[type_key] = field_path
                 nullable = False
+                primary_keys.append(field_path)
 
             field = SchemaField(
                 fieldPath=field_path,
                 nativeDataType=native_data_type,
                 type=type,
-                description=description,
+                description=None,
                 nullable=nullable,
                 recursive=False,
             )
@@ -513,6 +518,7 @@ class DynamoDBSource(StatefulIngestionSourceBase):
             hash="",
             platformSchema=SchemalessClass(),
             fields=canonical_schema,
+            primaryKeys=primary_keys,
         )
         return schema_metadata
 

datahub/ingestion/source/excel/config.py

@@ -0,0 +1,92 @@
+from typing import List, Optional, Union
+
+from pydantic.fields import Field
+
+from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.source_common import DatasetSourceConfigMixin
+from datahub.ingestion.source.aws.aws_common import AwsConnectionConfig
+from datahub.ingestion.source.azure.azure_common import AzureConnectionConfig
+from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StatefulStaleMetadataRemovalConfig,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionConfigBase,
+)
+from datahub.ingestion.source_config.operation_config import is_profiling_enabled
+
+
+class ExcelSourceConfig(StatefulIngestionConfigBase, DatasetSourceConfigMixin):
+    path_list: List[str] = Field(
+        description="List of paths to Excel files or folders to ingest."
+    )
+
+    path_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for file paths to filter in ingestion.",
+    )
+
+    aws_config: Optional[AwsConnectionConfig] = Field(
+        default=None, description="AWS configuration"
+    )
+
+    use_s3_bucket_tags: Optional[bool] = Field(
+        default=False,
+        description="Whether or not to create tags in datahub from the s3 bucket",
+    )
+
+    use_s3_object_tags: Optional[bool] = Field(
+        default=False,
+        description="Whether or not to create tags in datahub from the s3 object",
+    )
+
+    verify_ssl: Union[bool, str] = Field(
+        default=True,
+        description="Either a boolean, in which case it controls whether we verify the server's TLS certificate, or a string, in which case it must be a path to a CA bundle to use.",
+    )
+
+    azure_config: Optional[AzureConnectionConfig] = Field(
+        default=None, description="Azure configuration"
+    )
+
+    use_abs_blob_tags: Optional[bool] = Field(
+        default=False,
+        description="Whether to create tags in datahub from the abs blob tags",
+    )
+
+    convert_urns_to_lowercase: bool = Field(
+        default=False,
+        description="Enable to convert the Excel asset urns to lowercase",
+    )
+
+    active_sheet_only: bool = Field(
+        default=False,
+        description="Enable to only ingest the active sheet of the workbook. If not set, all sheets will be ingested.",
+    )
+
+    worksheet_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for worksheets to ingest. Worksheets are specified as 'filename_without_extension.worksheet_name'. "
+        "For example to allow the worksheet Sheet1 from file report.xlsx, use the pattern: 'report.Sheet1'.",
+    )
+
+    profile_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for worksheets to profile. Worksheets are specified as 'filename_without_extension.worksheet_name'. "
+        "For example to allow the worksheet Sheet1 from file report.xlsx, use the pattern: 'report.Sheet1'.",
+    )
+
+    profiling: GEProfilingConfig = Field(
+        default=GEProfilingConfig(),
+        description="Configuration for profiling",
+    )
+
+    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = Field(
+        default=None,
+        description="Configuration for stateful ingestion and stale metadata removal.",
+    )
+
+    def is_profiling_enabled(self) -> bool:
+        return self.profiling.enabled and is_profiling_enabled(
+            self.profiling.operation_config
+        )