acryl-datahub-cloud 0.3.10rc4__py3-none-any.whl → 0.3.16.1rc0__py3-none-any.whl

acryl_datahub_cloud/datahub_reporting/datahub_dataset.py

@@ -5,14 +5,14 @@ import pathlib
 import tempfile
 import time
 from enum import Enum
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union, cast
 
 import boto3
 import duckdb
 import pandas
 import pyarrow as pa
 import pyarrow.parquet as pq
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow, SchemaField
 from datahub.configuration.common import ConfigModel
@@ -73,7 +73,9 @@ class FileStoreBackedDatasetConfig(ConfigModel):
     store_platform: str = "s3"
     file_name: str = "data"
     file_extension: str = "parquet"
-    file_compression: str = "snappy"
+    file_compression: Literal[
+        "gzip", "bz2", "brotli", "lz4", "zstd", "snappy", "none"
+    ] = "snappy"
     file_overwrite_existing: bool = True
     snapshot_partitioning_strategy: str = PartitioningStrategy.DATE
     generate_presigned_url: bool = True
@@ -85,7 +87,8 @@ class FileStoreBackedDatasetConfig(ConfigModel):
 
     datahub_platform: str = "acryl"
 
-    @validator("snapshot_partitioning_strategy")
+    @field_validator("snapshot_partitioning_strategy")
+    @classmethod
     def validate_partitioning_strategy(cls, v):
         if v not in PartitioningStrategy._value2member_map_:
             raise ValueError(f"Unsupported partitioning strategy: {v}")
@@ -119,9 +122,14 @@ class DataHubBasedS3Dataset:
         self.local_file_path: str = (
             config.file if config.file else self._initialize_local_file()
         )
-        self.file_writer = None
+        self.file_writer: Optional[pq.ParquetWriter] = None
         self.schema = (
-            pa.schema([(x.name, x.type) for x in self.dataset_metadata.schemaFields])
+            pa.schema(
+                [
+                    pa.field(x.name, BaseModelRow.string_to_pyarrow_type(x.type))
+                    for x in self.dataset_metadata.schemaFields
+                ]
+            )
             if self.dataset_metadata.schemaFields
             else None
         )
@@ -163,18 +171,32 @@ class DataHubBasedS3Dataset:
                     self.schema = row.arrow_schema()
                 else:
                     # hail mary: infer schema from the first row and cast everything to string
-                    self.schema = pa.schema([(key, pa.string()) for key in row.keys()])
+                    self.schema = pa.schema([pa.field(key, pa.string()) for key in row])
                     self.stringify_row = True
 
             self._initialize_local_file()
+            # Map compression names to PyArrow format (most are direct mappings)
+            compression_map = {
+                "gzip": "gzip",
+                "bz2": "brotli",  # PyArrow doesn't support bz2, use brotli
+                "brotli": "brotli",
+                "lz4": "lz4",
+                "zstd": "zstd",
+                "snappy": "snappy",
+                "none": "none",
+            }
+            compression = cast(
+                Literal["gzip", "bz2", "brotli", "lz4", "zstd", "snappy", "none"],
+                compression_map.get(self.config.file_compression, "snappy"),
+            )
             self.file_writer = pq.ParquetWriter(
                 self.local_file_path,
                 self.schema,
-                compression=self.config.file_compression,
+                compression=compression,
             )
-        if isinstance(row, BaseModel) or isinstance(row, BaseModelRow):
+        if isinstance(row, (BaseModel, BaseModelRow)):
             # for anything extending BaseModel, we want to use the dict representation
-            write_row: Dict[str, Any] = row.dict()
+            write_row: Dict[str, Any] = row.model_dump()
         elif isinstance(row, dict):
             write_row = row
         else:
@@ -271,14 +293,10 @@ class DataHubBasedS3Dataset:
             )
 
     def _generate_schema_metadata(
-        self, duckdb_columns: List[Tuple[str, str]]
+        self, duckdb_columns: List[Tuple[str, Any]]
     ) -> SchemaMetadataClass:
         def get_type_from_dtype(dtype: str) -> SchemaFieldDataTypeClass:
-            if "int" in dtype:
-                return SchemaFieldDataTypeClass(type=NumberTypeClass())
-            elif "float" in dtype:
-                return SchemaFieldDataTypeClass(type=NumberTypeClass())
-            elif "number" in dtype:
+            if "int" in dtype or "float" in dtype or "number" in dtype:
                 return SchemaFieldDataTypeClass(type=NumberTypeClass())
             elif "bool" in dtype:
                 return SchemaFieldDataTypeClass(type=BooleanTypeClass())
@@ -306,7 +324,7 @@ class DataHubBasedS3Dataset:
         )
         for column in duckdb_columns:
             # generate data type
-            data_type = column[1].lower()
+            data_type = str(column[1]).lower()
             schema_metadata.fields.append(
                 SchemaFieldClass(
                     fieldPath=column[0],
@@ -345,7 +363,7 @@ class DataHubBasedS3Dataset:
                 # generate min, max, avg, distinct count, null count
                 column_name = column[0]
                 logger.info(f"Generating field profile for {column_name}")
-                data_type = column[1].lower()
+                data_type = str(column[1]).lower()
                 if "int" in data_type or "float" in data_type:
                     query = (
                         f"SELECT COUNT(DISTINCT {column_name}), COUNT(*) - COUNT({column_name}), MIN({column_name}), MAX({column_name}), AVG({column_name})"
@@ -400,7 +418,9 @@ class DataHubBasedS3Dataset:
                 assert dataset_profiles.fieldProfiles is not None
                 dataset_profiles.fieldProfiles.append(field_profile)
             logger.info("Generated dataset profile")
-            schema_metadata = self._generate_schema_metadata(columns)
+            schema_metadata = self._generate_schema_metadata(
+                [(col[0], col[1]) for col in columns]
+            )
         return dataset_profiles, schema_metadata
 
     def register_dataset(

acryl_datahub_cloud/datahub_reporting/datahub_form_reporting.py

@@ -1,13 +1,16 @@
-import json
 import logging
 from datetime import date, datetime, timezone
 from enum import Enum
-from typing import Any, Callable, Dict, Iterable, List, Optional
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 
 import pandas as pd
+from pydantic import BaseModel, field_validator
+
+from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow
+from acryl_datahub_cloud.graphql_utils import parse_extra_properties_for_model
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
-from datahub.ingestion.graph.filters import RawSearchFilterRule
+from datahub.ingestion.graph.filters import RawSearchFilter
 from datahub.metadata.schema_classes import (
     DomainPropertiesClass,
     FormAssociationClass,
@@ -16,9 +19,6 @@ from datahub.metadata.schema_classes import (
     FormStateClass,
     FormTypeClass,
 )
-from pydantic import BaseModel
-
-from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow
 
 logger = logging.getLogger(__name__)
 
@@ -130,6 +130,22 @@ class DataHubFormReportingData(FormData):
         platformInstance: Optional[str] = None
         domains: List[str] = []
 
+        @field_validator(
+            "completedFormsIncompletePromptResponseTimes",
+            "completedFormsCompletedPromptResponseTimes",
+            "incompleteFormsIncompletePromptResponseTimes",
+            "incompleteFormsCompletedPromptResponseTimes",
+            mode="before",
+        )
+        @classmethod
+        def convert_timestamps_to_strings(
+            cls, v: Union[List[int], List[str]]
+        ) -> List[str]:
+            """Convert timestamp integers to strings for compatibility with GMS data."""
+            if not isinstance(v, list):
+                return v
+            return [str(item) for item in v]
+
     def __init__(self, graph: DataHubGraph, allowed_forms: Optional[List[str]] = None):
         self.graph: DataHubGraph = graph
         self.form_registry = FormRegistry(graph)
@@ -143,13 +159,13 @@ class DataHubFormReportingData(FormData):
         on_form_scanned: Callable[[str], Any],
     ) -> pd.DataFrame:
         return pd.DataFrame(
-            x.dict()
+            x.model_dump()
             for x in self.get_data(
                 on_asset_scanned=on_asset_scanned, on_form_scanned=on_form_scanned
             )
         )
 
-    def get_form_existence_or_filters(self) -> List[RawSearchFilterRule]:
+    def get_form_existence_or_filters(self) -> RawSearchFilter:
         """
         Datasets must either have completedForms or incompleteForms assigned to
         them
@@ -157,25 +173,41 @@ class DataHubFormReportingData(FormData):
         if self.allowed_forms:
             return [
                 {
-                    "field": "completedForms",
-                    "condition": "EQUAL",
-                    "values": self.allowed_forms,
+                    "and": [
+                        {
+                            "field": "completedForms",
+                            "condition": "EQUAL",
+                            "values": self.allowed_forms,
+                        }
+                    ]
                 },
                 {
-                    "field": "incompleteForms",
-                    "condition": "EQUAL",
-                    "values": self.allowed_forms,
+                    "and": [
+                        {
+                            "field": "incompleteForms",
+                            "condition": "EQUAL",
+                            "values": self.allowed_forms,
+                        }
+                    ]
                 },
             ]
         else:
             return [
                 {
-                    "field": "completedForms",
-                    "condition": "EXISTS",
+                    "and": [
+                        {
+                            "field": "completedForms",
+                            "condition": "EXISTS",
+                        }
+                    ]
                 },
                 {
-                    "field": "incompleteForms",
-                    "condition": "EXISTS",
+                    "and": [
+                        {
+                            "field": "incompleteForms",
+                            "condition": "EXISTS",
+                        }
+                    ]
                 },
             ]
 
@@ -257,6 +289,7 @@ class DataHubFormReportingData(FormData):
                 for prompt_id, response_time in zip(
                     search_row.completedFormsCompletedPromptIds,
                     search_row.completedFormsCompletedPromptResponseTimes,
+                    strict=False,
                 )
                 if prompt_id in form_prompts
             }
@@ -289,7 +322,8 @@ class DataHubFormReportingData(FormData):
         on_asset_scanned: Optional[Callable[[str], Any]] = None,
         on_form_scanned: Optional[Callable[[str], Any]] = None,
     ) -> Iterable[FormReportingRow]:
-        extra_fields = [f for f in self.DataHubDatasetSearchRow.__fields__.keys()]
+        extra_fields = [f for f in self.DataHubDatasetSearchRow.model_fields]
+        # TODO: Replace with the new search/filter SDK.
         result = self.graph.get_results_by_filter(
             extra_or_filters=self.get_form_existence_or_filters(),
             extra_source_fields=extra_fields,
@@ -302,10 +336,9 @@ class DataHubFormReportingData(FormData):
             if row_index % 100 == 0:
                 logger.info(f"Scanned {row_index} assets")
             extra_properties = row["extraProperties"]
-
-            extra_properties_map = {
-                x["name"]: json.loads(x["value"]) for x in extra_properties
-            }
+            extra_properties_map = parse_extra_properties_for_model(
+                extra_properties, self.DataHubDatasetSearchRow
+            )
             search_row = self.DataHubDatasetSearchRow(**extra_properties_map)
             if on_asset_scanned:
                 on_asset_scanned(search_row.urn)
@@ -388,6 +421,7 @@ class DataHubFormReportingData(FormData):
                     for (p, p_response_time) in zip(
                         search_row.incompleteFormsCompletedPromptIds,
                         search_row.incompleteFormsCompletedPromptResponseTimes,
+                        strict=False,
                     )
                     if p in form_prompts
                 ]:
@@ -411,7 +445,7 @@ class DataHubFormReportingData(FormData):
                             question_status=QuestionStatus.COMPLETED,
                             question_completed_date=datetime.fromtimestamp(
                                 float(prompt_response_time) / 1000, tz=timezone.utc
-                            ),
+                            ).date(),
                             snapshot_date=self.snapshot_date,
                         )
             complete_forms = (
@@ -485,6 +519,7 @@ class DataHubFormReportingData(FormData):
                     for (p, p_response_time) in zip(
                         search_row.completedFormsCompletedPromptIds,
                         search_row.completedFormsCompletedPromptResponseTimes,
+                        strict=False,
                     )
                     if p in form_prompts
                 ]:
@@ -512,7 +547,7 @@ class DataHubFormReportingData(FormData):
                             question_status=QuestionStatus.COMPLETED,
                             question_completed_date=datetime.fromtimestamp(
                                 float(prompt_response_time) / 1000, tz=timezone.utc
-                            ),
+                            ).date(),
                             snapshot_date=self.snapshot_date,
                         )
 

acryl_datahub_cloud/datahub_reporting/extract_graph.py

@@ -6,7 +6,7 @@ from typing import List, Optional
 
 import boto3
 from opensearchpy import OpenSearch
-from pydantic import validator
+from pydantic import field_validator
 
 from acryl_datahub_cloud.datahub_reporting.datahub_dataset import (
     DataHubBasedS3Dataset,
@@ -32,6 +32,7 @@ logger = logging.getLogger(__name__)
 
 
 class DataHubReportingExtractGraphSourceConfig(ConfigModel):
+    enabled: bool = True
     server: Optional[DatahubClientConfig] = None
     search_index: ElasticSearchClientConfig = ElasticSearchClientConfig()
     extract_graph_store: FileStoreBackedDatasetConfig
@@ -42,8 +43,9 @@ class DataHubReportingExtractGraphSourceConfig(ConfigModel):
     query_timeout: int = 30
     extract_batch_size: int = 2000
 
-    @validator("extract_graph_store", pre=True, always=True)
-    def set_default_extract_soft_delete_flag(cls, v, values):
+    @field_validator("extract_graph_store", mode="before")
+    @classmethod
+    def set_default_extract_soft_delete_flag(cls, v):
         if v is not None:
             if "dataset_registration_spec" not in v:
                 v["dataset_registration_spec"] = DatasetRegistrationSpec(
@@ -118,6 +120,10 @@ class DataHubReportingExtractGraphSource(Source):
         return skip_extract
 
     def get_workunits(self):
+        if not self.config.enabled:
+            logger.info("Source is disabled, stopping")
+            return
+
         self.graph = (
             self.ctx.require_graph("Loading default graph coordinates.")
             if self.config.server is None