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

acryl_datahub_cloud/elasticsearch/graph_service.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, get_args, get_origin
 
 import pyarrow as pa
 from pydantic import BaseModel
@@ -15,7 +15,37 @@ class BaseModelRow(BaseModel):
         return self.__dict__
 
     @staticmethod
-    def pydantic_type_to_pyarrow(type_: Any) -> pa.DataType:
+    def pydantic_type_to_pyarrow(type_: Any) -> pa.DataType:  # noqa: C901
+        # Handle generic types (List, Optional, Union, etc.)
+        origin = get_origin(type_)
+
+        if origin is list:
+            # List[X] -> list(X)
+            args = get_args(type_)
+            if args:
+                inner_type = BaseModelRow.pydantic_type_to_pyarrow(args[0])
+                return pa.list_(inner_type)
+            return pa.list_(pa.string())  # Default to list of strings
+
+        if origin is type(None):
+            # Just None type
+            return pa.null()
+
+        # Optional[X] is Union[X, None]
+        args = get_args(type_)
+        if args:
+            # Check if this is Optional (Union with None)
+            non_none_types = [arg for arg in args if arg is not type(None)]
+            if len(non_none_types) < len(args):  # Had None in the union
+                if non_none_types:
+                    return BaseModelRow.pydantic_type_to_pyarrow(non_none_types[0])
+                return pa.null()
+
+        # Handle simple types - check if it's actually a class first
+        if not isinstance(type_, type):
+            # If it's not a type/class, default to string
+            return pa.string()
+
         if issubclass(type_, bool):
             return pa.bool_()
         elif issubclass(type_, int):
@@ -32,28 +62,60 @@ class BaseModelRow(BaseModel):
         else:
             raise ValueError(f"No mapping for type {type_}")
 
+    @staticmethod
+    def string_to_pyarrow_type(type_string: str) -> pa.DataType:
+        """Convert string representation back to pyarrow type by converting to Python type first."""
+        # Mapping of pyarrow string representations to Python types
+        type_mapping = {
+            "string": str,
+            "int64": int,
+            "float64": float,
+            "bool": bool,
+            "timestamp[ns]": datetime.datetime,
+            "date32[day]": datetime.date,
+        }
+
+        python_type = type_mapping.get(
+            type_string, str
+        )  # Default to str for unknown types
+        return BaseModelRow.pydantic_type_to_pyarrow(python_type)
+
+    @staticmethod
+    def _is_optional_type(type_: Any) -> bool:
+        """Check if a type annotation is Optional (i.e., Union with None)."""
+        origin = get_origin(type_)
+        if origin is type(None):
+            return True
+
+        args = get_args(type_)
+        if args:
+            # Check if None is in the Union args (Optional is Union[X, None])
+            return type(None) in args
+
+        return False
+
     @classmethod
     def arrow_schema(cls) -> pa.Schema:
         fields = []
-        for field_name, field_model in cls.__fields__.items():
-            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(
-                field_model.outer_type_
-            )
+        for field_name, field_model in cls.model_fields.items():
+            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(field_model.annotation)
             pyarrow_field = pa.field(field_name, pyarrow_type)
-            if not field_model.required:
-                pyarrow_field = pyarrow_field.with_nullable(True)
-            else:
-                pyarrow_field = pyarrow_field.with_nullable(False)
+
+            # Check if the type is Optional (Union with None) OR if the field has a default value
+            is_nullable = (
+                BaseModelRow._is_optional_type(field_model.annotation)
+                or not field_model.is_required()
+            )
+
+            pyarrow_field = pyarrow_field.with_nullable(is_nullable)
             fields.append(pyarrow_field)
         return pa.schema(fields)
 
     @classmethod
     def datahub_schema(cls) -> List[SchemaField]:
         fields = []
-        for field_name, field_model in cls.__fields__.items():
-            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(
-                field_model.outer_type_
-            )
+        for field_name, field_model in cls.model_fields.items():
+            pyarrow_type = BaseModelRow.pydantic_type_to_pyarrow(field_model.annotation)
             fields.append(SchemaField(name=field_name, type=str(pyarrow_type)))
         return fields
 

acryl_datahub_cloud/graphql_utils.py

@@ -0,0 +1,64 @@
+import json
+import logging
+from typing import Any, Dict, List, Type, TypeVar
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def parse_extra_properties_for_model(
+    extra_properties: List[Dict[str, str]], model_class: Type[T]
+) -> Dict[str, Any]:
+    """
+    Parse GraphQL extraProperties into a dictionary suitable for Pydantic model instantiation.
+
+    GraphQL returns extra properties where values are JSON-encoded strings. This function:
+    1. Filters to only include fields defined in the model
+    2. Parses JSON-encoded values back to Python objects
+    3. Handles parsing errors gracefully with logging
+
+    Args:
+        extra_properties: List of dicts with 'name' and 'value' keys from GraphQL response
+        model_class: Pydantic model class to parse properties for
+
+    Returns:
+        Dictionary mapping field names to parsed values, ready for model instantiation
+
+    Example:
+        extra_props = [
+            {"name": "urn", "value": '"urn:li:dataset:..."'},
+            {"name": "owners", "value": '["urn:li:corpuser:user1"]'},
+            {"name": "scrollId", "value": "eyJzb3J0Ijp..."}  # Not in model, will be skipped
+        ]
+        parsed = parse_extra_properties_for_model(extra_props, MyModel)
+        instance = MyModel(**parsed)
+    """
+    # Only parse fields that are defined in the model.
+    # Extra fields like scrollId are ignored.
+    model_fields = model_class.model_fields.keys()
+    extra_properties_map: Dict[str, Any] = {}
+
+    for prop in extra_properties:
+        name = prop["name"]
+        # Skip fields that aren't in the model (e.g., scrollId)
+        if name not in model_fields:
+            continue
+
+        # All model fields should be JSON-encoded, parse them
+        value = prop["value"]
+        try:
+            extra_properties_map[name] = json.loads(value)
+        except json.JSONDecodeError as e:
+            # This indicates a data contract violation - the field should be JSON-encoded
+            logger.warning(
+                f"Failed to parse field '{name}' as JSON. "
+                f"Value: '{value[:100]}...'. Error: {e}. "
+                f"Skipping field and using default value."
+            )
+            # Skip this field - Pydantic will use the default value if available,
+            # or fail validation if it's a required field
+
+    return extra_properties_map