Merge pull request #396 from atlanhq/DVX-657

DVX-657: Added experimental support for serde on the `DataContract` spec

pyatlan/client/contract.py

@@ -5,6 +5,7 @@
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import CONTRACT_INIT_API
 from pyatlan.errors import ErrorCode
+from pyatlan.model.assets import Asset
 from pyatlan.model.contract import InitRequest
 
 
@@ -22,20 +23,22 @@ def __init__(self, client: ApiCaller):
 
     @validate_arguments
     def generate_initial_spec(
-        self, asset_type: str, asset_qualified_name: str
+        self,
+        asset: Asset,
     ) -> Optional[str]:
         """
-        Generate an initial contract spec for the provided asset type and qualified name.
+        Generate an initial contract spec for the provided asset.
+        The asset must have at least its `qualifiedName` (and `typeName`) populated.
 
-        :param asset_type: `typeName` of the asset, eg: `Table`, `Column` etc
-        :param asset_qualified_name: `qualifiedName` of the asset
-        :raises AtlanError: on any issue interacting with the API
-        :returns: `YAML` for the initial contract spec for the provided asset
+        :param asset: for which to generate the initial contract spec
+
+        :raises AtlanError: if there is an issue interacting with the API
+        :returns: YAML for the initial contract spec for the provided asset
         """
         response = self._client._call_api(
             CONTRACT_INIT_API,
             request_obj=InitRequest(
-                asset_type=asset_type, asset_qualified_name=asset_qualified_name
+                asset_type=asset.type_name, asset_qualified_name=asset.qualified_name
             ),
         )
         return response.get("contract")

pyatlan/generator/templates/methods/asset/data_contract.jinja2

@@ -1,13 +1,30 @@
 
+    @overload
+    @classmethod
+    def creator(cls, asset_qualified_name: str, contract_json: str) -> DataContract: ...
+
+    @overload
+    @classmethod
+    def creator(
+        cls, asset_qualified_name: str, contract_spec: Union[DataContractSpec, str]
+    ) -> DataContract: ...
+
     @classmethod
     @init_guid
-    def creator(cls, *, asset_qualified_name: str, contract_json: str) -> DataContract:
+    def creator(
+        cls,
+        *,
+        asset_qualified_name: str,
+        contract_json: Optional[str] = None,
+        contract_spec: Optional[Union[DataContractSpec, str]] = None,
+    ) -> DataContract:
         validate_required_fields(
-            ["asset_qualified_name", "contract_json"],
-            [asset_qualified_name, contract_json],
+            ["asset_qualified_name"],
+            [asset_qualified_name],
         )
         attributes = DataContract.Attributes.creator(
             asset_qualified_name=asset_qualified_name,
             contract_json=contract_json,
+            contract_spec=contract_spec,
         )
         return cls(attributes=attributes)

pyatlan/generator/templates/methods/attribute/data_contract.jinja2

@@ -2,19 +2,55 @@
         @classmethod
         @init_guid
         def creator(
-            cls, *, asset_qualified_name: str, contract_json: str
+            cls,
+            *,
+            asset_qualified_name: str,
+            contract_json: Optional[str] = None,
+            contract_spec: Optional[Union[DataContractSpec, str]] = None,
         ) -> DataContract.Attributes:
+            from re import search
+
             validate_required_fields(
-                ["asset_qualified_name", "contract_json"],
-                [asset_qualified_name, contract_json],
+                ["asset_qualified_name"],
+                [asset_qualified_name],
             )
-            try:
-                contract_name = f"Data contract for {loads(contract_json)['dataset']}"
-            except (JSONDecodeError, KeyError):
-                raise ErrorCode.INVALID_CONTRACT_JSON.exception_with_parameters()
+            if not (contract_json or contract_spec):
+                raise ValueError(
+                    "At least one of `contract_json` or `contract_spec` "
+                    "must be provided to create a contract."
+                )
+            if contract_json and contract_spec:
+                raise ValueError(
+                    "Both `contract_json` and `contract_spec` cannot be "
+                    "provided simultaneously to create a contract."
+                )
+            last_slash_index = asset_qualified_name.rfind("/")
+            default_dataset = asset_qualified_name[last_slash_index + 1 :]  # noqa
+
+            if contract_json:
+                try:
+                    contract_name = f"Data contract for {loads(contract_json)['dataset'] or default_dataset}"
+                except (JSONDecodeError, KeyError):
+                    raise ErrorCode.INVALID_CONTRACT_JSON.exception_with_parameters()
+            else:
+                if isinstance(contract_spec, DataContractSpec):
+                    contract_name = (
+                        "Data contract for "
+                        f"{contract_spec.dataset or default_dataset}"  # type: ignore[union-attr, attr-defined]
+                    )
+                    contract_spec = contract_spec.to_yaml()
+                else:
+                    is_dataset_found = search(
+                        r"dataset:\s*([^\s#]+)", contract_spec or default_dataset
+                    )
+                    dataset = None
+                    if is_dataset_found:
+                        dataset = is_dataset_found.group(1)
+                    contract_name = f"Data contract for {dataset or default_dataset}"
 
             return DataContract.Attributes(
                 name=contract_name,
                 qualified_name=f"{asset_qualified_name}/contract",
                 data_contract_json=contract_json,
+                data_contract_spec=contract_spec,  # type: ignore[arg-type]
             )

pyatlan/model/assets/core/data_contract.py

@@ -6,11 +6,12 @@
 
 from json import loads
 from json.decoder import JSONDecodeError
-from typing import ClassVar, List, Optional
+from typing import ClassVar, List, Optional, Union, overload
 
 from pydantic.v1 import Field, validator
 
 from pyatlan.errors import ErrorCode
+from pyatlan.model.contract import DataContractSpec
 from pyatlan.model.fields.atlan_fields import (
     KeywordField,
     NumericField,
@@ -25,16 +26,33 @@
 class DataContract(Catalog):
     """Description"""
 
+    @overload
+    @classmethod
+    def creator(cls, asset_qualified_name: str, contract_json: str) -> DataContract: ...
+
+    @overload
+    @classmethod
+    def creator(
+        cls, asset_qualified_name: str, contract_spec: Union[DataContractSpec, str]
+    ) -> DataContract: ...
+
     @classmethod
     @init_guid
-    def creator(cls, *, asset_qualified_name: str, contract_json: str) -> DataContract:
+    def creator(
+        cls,
+        *,
+        asset_qualified_name: str,
+        contract_json: Optional[str] = None,
+        contract_spec: Optional[Union[DataContractSpec, str]] = None,
+    ) -> DataContract:
         validate_required_fields(
-            ["asset_qualified_name", "contract_json"],
-            [asset_qualified_name, contract_json],
+            ["asset_qualified_name"],
+            [asset_qualified_name],
         )
         attributes = DataContract.Attributes.creator(
             asset_qualified_name=asset_qualified_name,
             contract_json=contract_json,
+            contract_spec=contract_spec,
         )
         return cls(attributes=attributes)
 
@@ -241,21 +259,57 @@ class Attributes(Catalog.Attributes):
         @classmethod
         @init_guid
         def creator(
-            cls, *, asset_qualified_name: str, contract_json: str
+            cls,
+            *,
+            asset_qualified_name: str,
+            contract_json: Optional[str] = None,
+            contract_spec: Optional[Union[DataContractSpec, str]] = None,
         ) -> DataContract.Attributes:
+            from re import search
+
             validate_required_fields(
-                ["asset_qualified_name", "contract_json"],
-                [asset_qualified_name, contract_json],
+                ["asset_qualified_name"],
+                [asset_qualified_name],
             )
-            try:
-                contract_name = f"Data contract for {loads(contract_json)['dataset']}"
-            except (JSONDecodeError, KeyError):
-                raise ErrorCode.INVALID_CONTRACT_JSON.exception_with_parameters()
+            if not (contract_json or contract_spec):
+                raise ValueError(
+                    "At least one of `contract_json` or `contract_spec` "
+                    "must be provided to create a contract."
+                )
+            if contract_json and contract_spec:
+                raise ValueError(
+                    "Both `contract_json` and `contract_spec` cannot be "
+                    "provided simultaneously to create a contract."
+                )
+            last_slash_index = asset_qualified_name.rfind("/")
+            default_dataset = asset_qualified_name[last_slash_index + 1 :]  # noqa
+
+            if contract_json:
+                try:
+                    contract_name = f"Data contract for {loads(contract_json)['dataset'] or default_dataset}"
+                except (JSONDecodeError, KeyError):
+                    raise ErrorCode.INVALID_CONTRACT_JSON.exception_with_parameters()
+            else:
+                if isinstance(contract_spec, DataContractSpec):
+                    contract_name = (
+                        "Data contract for "
+                        f"{contract_spec.dataset or default_dataset}"  # type: ignore[union-attr, attr-defined]
+                    )
+                    contract_spec = contract_spec.to_yaml()
+                else:
+                    is_dataset_found = search(
+                        r"dataset:\s*([^\s#]+)", contract_spec or default_dataset
+                    )
+                    dataset = None
+                    if is_dataset_found:
+                        dataset = is_dataset_found.group(1)
+                    contract_name = f"Data contract for {dataset or default_dataset}"
 
             return DataContract.Attributes(
                 name=contract_name,
                 qualified_name=f"{asset_qualified_name}/contract",
                 data_contract_json=contract_json,
+                data_contract_spec=contract_spec,  # type: ignore[arg-type]
             )
 
     attributes: DataContract.Attributes = Field(