Merge pull request #446 from atlanhq/FT-617

FT-617: Added exposure of source-specific custom attributes

pyatlan/model/assets/core/referenceable.py

@@ -4,9 +4,10 @@
 
 from __future__ import annotations
 
+from json import JSONDecodeError, loads
 from typing import Any, ClassVar, Dict, List, Optional
 
-from pydantic.v1 import Field, PrivateAttr
+from pydantic.v1 import Field, PrivateAttr, root_validator
 
 from pyatlan.model.core import AtlanObject, AtlanTag, Meaning
 from pyatlan.model.custom_metadata import CustomMetadataDict, CustomMetadataProxy
@@ -18,6 +19,7 @@
     KeywordField,
     KeywordTextField,
     NumericField,
+    TextField,
 )
 from pyatlan.model.lineage_ref import LineageRef
 
@@ -32,6 +34,23 @@ def __init__(__pydantic_self__, **data: Any) -> None:
             __pydantic_self__.business_attributes
         )
 
+    @root_validator(pre=True)
+    def parse_custom_attributes(cls, values):
+        if "attributes" in values:
+            attributes = values["attributes"]
+            if "__customAttributes" in attributes:
+                # Pop the __customAttributes from attributes
+                custom_attributes = attributes.pop("__customAttributes")
+                try:
+                    # Try to parse the JSON string if it's a string
+                    if isinstance(custom_attributes, str):
+                        custom_attributes = loads(custom_attributes)
+                    # Add the parsed custom attributes to the Column
+                    values["custom_attributes"] = custom_attributes
+                except JSONDecodeError:
+                    pass
+        return values
+
     def json(self, *args, **kwargs) -> str:
         self.business_attributes = self._metadata_proxy.business_attributes
         return super().json(**kwargs)
@@ -222,6 +241,15 @@ def validate_required(self):
     )
     """Unique fully-qualified name of the asset in Atlan."""
 
+    CUSTOM_ATTRIBUTES: ClassVar[TextField] = TextField(
+        "__customAttributes", "__customAttributes"
+    )
+    """
+    Any source-provided custom information.
+    NOTE: This is NOT the same as custom metadata (user-managed),
+    but is an entirely different area of source-managed custom information.
+    """
+
     type_name: str = Field(
         default="Referenceable",
         description="Name of the type definition that defines this instance.",

tests/integration/test_index_search.py

@@ -37,6 +37,10 @@
 NOW_AS_YYYY_MM_DD = datetime.today().strftime("%Y-%m-%d")
 EXISTING_TAG = "Issue"
 EXISTING_SOURCE_SYNCED_TAG = "Confidential"
+COLUMN_WITH_CUSTOM_ATTRIBUTES = (
+    "default/snowflake/1733440936/ANALYTICS"
+    "/WIDE_WORLD_IMPORTERS/STG_STATE_PROVINCES/LATEST_RECORDED_POPULATION"
+)
 
 VALUES_FOR_TERM_QUERIES = {
     "with_categories": "VBsYc9dUoEcAtDxZmjby6@mweSfpXBwfYWedQTvA3Gi",
@@ -243,6 +247,31 @@ def test_source_tag_assign_with_value(client: AtlanClient, table: Table):
     _assert_source_tag(tables, EXISTING_SOURCE_SYNCED_TAG, "Not Restricted")
 
 
+def test_search_source_specific_custom_attributes(client: AtlanClient):
+    # Test with get_by_qualified_name()
+    column_qn = COLUMN_WITH_CUSTOM_ATTRIBUTES
+    asset = client.asset.get_by_qualified_name(
+        asset_type=Column,
+        qualified_name=column_qn,
+        min_ext_info=True,
+        ignore_relationships=True,
+    )
+    assert asset and asset.custom_attributes
+
+    # Test with FluentSearch()
+    results = (
+        FluentSearch()
+        .where(CompoundQuery.active_assets())
+        .where(Column.QUALIFIED_NAME.eq(column_qn))
+        .include_on_results(Column.CUSTOM_ATTRIBUTES)
+        .execute(client=client)
+    )
+    assert results and results.count == 1
+    assert results.current_page() and len(results.current_page()) == 1
+    column = results.current_page()[0]
+    assert isinstance(column, Column) and column and column.custom_attributes
+
+
 def test_search_next_page(client: AtlanClient):
     size = 2
     dsl = DSL(