Feat: Refactor schema classes to subclass from _Resource (#1536)

* feat: Refactor code to make shema type classes a subclass of Artifact, Execution and Context

* reinstantiate schema class using resouce_name

* fix typo in resource_name

* change resouce_name to artifact_name

* Change constructor to internal method

* refactor execution.

* refactor context.

* fix spelling

* fix unit tests

* fix sample tests.

Author: SinaChavoshi

google/cloud/aiplatform/base.py

@@ -727,7 +727,7 @@ def __repr__(self) -> str:
 
     def to_dict(self) -> Dict[str, Any]:
         """Returns the resource proto as a dictionary."""
-        return json_format.MessageToDict(self.gca_resource._pb)
+        return json_format.MessageToDict(self._gca_resource._pb)
 
     @classmethod
     def _generate_display_name(cls, prefix: Optional[str] = None) -> str:

google/cloud/aiplatform/metadata/artifact.py

@@ -31,7 +31,6 @@
 from google.cloud.aiplatform.metadata import metadata_store
 from google.cloud.aiplatform.metadata import resource
 from google.cloud.aiplatform.metadata import utils as metadata_utils
-from google.cloud.aiplatform.metadata.schema import base_artifact
 from google.cloud.aiplatform.utils import rest_utils
 
 
@@ -327,60 +326,15 @@ def create(
             credentials=credentials,
         )
 
-    @classmethod
-    def create_from_base_artifact_schema(
-        cls,
-        *,
-        base_artifact_schema: "base_artifact.BaseArtifactSchema",
-        metadata_store_id: Optional[str] = "default",
-        project: Optional[str] = None,
-        location: Optional[str] = None,
-        credentials: Optional[auth_credentials.Credentials] = None,
-    ) -> "Artifact":
-        """Creates a new Metadata Artifact from a BaseArtifactSchema class instance.
-
-        Args:
-            base_artifact_schema (BaseArtifactSchema):
-                Required. An instance of the BaseArtifactType class that can be
-                provided instead of providing artifact specific parameters.
-            metadata_store_id (str):
-                Optional. The <metadata_store_id> portion of the resource name with
-                the format:
-                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>
-                If not provided, the MetadataStore's ID will be set to "default".
-            project (str):
-                Optional. Project used to create this Artifact. Overrides project set in
-                aiplatform.init.
-            location (str):
-                Optional. Location used to create this Artifact. Overrides location set in
-                aiplatform.init.
-            credentials (auth_credentials.Credentials):
-                Optional. Custom credentials used to create this Artifact. Overrides
-                credentials set in aiplatform.init.
-
-        Returns:
-            Artifact: Instantiated representation of the managed Metadata Artifact.
-        """
-
-        return cls.create(
-            resource_id=base_artifact_schema.artifact_id,
-            schema_title=base_artifact_schema.schema_title,
-            uri=base_artifact_schema.uri,
-            display_name=base_artifact_schema.display_name,
-            schema_version=base_artifact_schema.schema_version,
-            description=base_artifact_schema.description,
-            metadata=base_artifact_schema.metadata,
-            state=base_artifact_schema.state,
-            metadata_store_id=metadata_store_id,
-            project=project,
-            location=location,
-            credentials=credentials,
-        )
-
     @property
     def uri(self) -> Optional[str]:
         "Uri for this Artifact."
-        return self.gca_resource.uri
+        return self._gca_resource.uri
+
+    @property
+    def state(self) -> Optional[gca_artifact.Artifact.State]:
+        "The State for this Artifact."
+        return self._gca_resource.state
 
     @classmethod
     def get_with_uri(

google/cloud/aiplatform/metadata/execution.py

@@ -31,7 +31,6 @@
 from google.cloud.aiplatform.metadata import artifact
 from google.cloud.aiplatform.metadata import metadata_store
 from google.cloud.aiplatform.metadata import resource
-from google.cloud.aiplatform.metadata.schema import base_execution
 
 
 class Execution(resource._Resource):
@@ -167,57 +166,6 @@ def create(
 
         return self
 
-    @classmethod
-    def create_from_base_execution_schema(
-        cls,
-        *,
-        base_execution_schema: "base_execution.BaseExecutionSchema",
-        metadata_store_id: Optional[str] = "default",
-        project: Optional[str] = None,
-        location: Optional[str] = None,
-        credentials: Optional[auth_credentials.Credentials] = None,
-    ) -> "Execution":
-        """
-        Creates a new Metadata Execution.
-
-        Args:
-            base_execution_schema (BaseExecutionSchema):
-                An instance of the BaseExecutionSchema class that can be
-                provided instead of providing schema specific parameters.
-            metadata_store_id (str):
-                Optional. The <metadata_store_id> portion of the resource name with
-                the format:
-                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>
-                If not provided, the MetadataStore's ID will be set to "default".
-            project (str):
-                Optional. Project used to create this Execution. Overrides project set in
-                aiplatform.init.
-            location (str):
-                Optional. Location used to create this Execution. Overrides location set in
-                aiplatform.init.
-            credentials (auth_credentials.Credentials):
-                Optional. Custom credentials used to create this Execution. Overrides
-                credentials set in aiplatform.init.
-
-        Returns:
-            Execution: Instantiated representation of the managed Metadata Execution.
-
-        """
-        resource = Execution.create(
-            state=base_execution_schema.state,
-            schema_title=base_execution_schema.schema_title,
-            resource_id=base_execution_schema.execution_id,
-            display_name=base_execution_schema.display_name,
-            schema_version=base_execution_schema.schema_version,
-            metadata=base_execution_schema.metadata,
-            description=base_execution_schema.description,
-            metadata_store_id=metadata_store_id,
-            project=project,
-            location=location,
-            credentials=credentials,
-        )
-        return resource
-
     def __enter__(self):
         if self.state is not gca_execution.Execution.State.RUNNING:
             self.update(state=gca_execution.Execution.State.RUNNING)

google/cloud/aiplatform/metadata/resource.py

@@ -114,6 +114,14 @@ def schema_title(self) -> str:
     def description(self) -> str:
         return self._gca_resource.description
 
+    @property
+    def display_name(self) -> str:
+        return self._gca_resource.display_name
+
+    @property
+    def schema_version(self) -> str:
+        return self._gca_resource.schema_version
+
     @classmethod
     def get_or_create(
         cls,
@@ -221,7 +229,7 @@ def get(
 
         Returns:
             resource (_Resource):
-                Instantiated representation of the managed Metadata resource or None if no resouce was found.
+                Instantiated representation of the managed Metadata resource or None if no resource was found.
 
         """
         resource = cls._get(

google/cloud/aiplatform/metadata/schema/base_artifact.py

@@ -26,7 +26,7 @@
 from google.cloud.aiplatform.metadata import constants
 
 
-class BaseArtifactSchema(metaclass=abc.ABCMeta):
+class BaseArtifactSchema(artifact.Artifact):
     """Base class for Metadata Artifact types."""
 
     @property
@@ -81,13 +81,40 @@ def __init__(
                 Pipelines), and the system does not prescribe or
                 check the validity of state transitions.
         """
+        # resource_id is not stored in the proto. Create method uses the
+        # resource_id along with project_id and location to construct an
+        # resource_name which is stored in the proto message.
         self.artifact_id = artifact_id
-        self.uri = uri
-        self.display_name = display_name
-        self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION
-        self.description = description
-        self.metadata = metadata
-        self.state = state
+
+        # Store all other attributes using the proto structure.
+        self._gca_resource = gca_artifact.Artifact()
+        self._gca_resource.uri = uri
+        self._gca_resource.display_name = display_name
+        self._gca_resource.schema_version = (
+            schema_version or constants._DEFAULT_SCHEMA_VERSION
+        )
+        self._gca_resource.description = description
+
+        # If metadata is None covert to {}
+        metadata = metadata if metadata else {}
+        self._nested_update_metadata(self._gca_resource, metadata)
+        self._gca_resource.state = state
+
+    # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
+    def _init_with_resource_name(
+        self,
+        *,
+        artifact_name: str,
+    ):
+
+        """Initializes the Artifact instance using an existing resource.
+
+        Args:
+            artifact_name (str):
+                Artifact name with the following format, this is globally unique in a metadataStore:
+                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
+        """
+        super(BaseArtifactSchema, self).__init__(artifact_name=artifact_name)
 
     def create(
         self,
@@ -117,10 +144,27 @@ def create(
         Returns:
             Artifact: Instantiated representation of the managed Metadata Artifact.
         """
-        return artifact.Artifact.create_from_base_artifact_schema(
-            base_artifact_schema=self,
+
+        # Check if metadata exists to avoid proto read error
+        metadata = None
+        if self._gca_resource.metadata:
+            metadata = self.metadata
+
+        new_artifact_instance = artifact.Artifact.create(
+            resource_id=self.artifact_id,
+            schema_title=self.schema_title,
+            uri=self.uri,
+            display_name=self.display_name,
+            schema_version=self.schema_version,
+            description=self.description,
+            metadata=metadata,
+            state=self.state,
             metadata_store_id=metadata_store_id,
             project=project,
             location=location,
             credentials=credentials,
         )
+
+        # Reinstantiate this class using the newly created resource.
+        self._init_with_resource_name(artifact_name=new_artifact_instance.resource_name)
+        return self

google/cloud/aiplatform/metadata/schema/base_context.py

@@ -21,11 +21,12 @@
 
 from google.auth import credentials as auth_credentials
 
+from google.cloud.aiplatform.compat.types import context as gca_context
 from google.cloud.aiplatform.metadata import constants
 from google.cloud.aiplatform.metadata import context
 
 
-class BaseContextSchema(metaclass=abc.ABCMeta):
+class BaseContextSchema(context.Context):
     """Base class for Metadata Context schema."""
 
     @property
@@ -62,11 +63,35 @@ def __init__(
             description (str):
                 Optional. Describes the purpose of the Context to be created.
         """
+        # resource_id is not stored in the proto. Create method uses the
+        # resource_id along with project_id and location to construct an
+        # resource_name which is stored in the proto message.
         self.context_id = context_id
-        self.display_name = display_name
-        self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION
-        self.metadata = metadata
-        self.description = description
+
+        # Store all other attributes using the proto structure.
+        self._gca_resource = gca_context.Context()
+        self._gca_resource.display_name = display_name
+        self._gca_resource.schema_version = (
+            schema_version or constants._DEFAULT_SCHEMA_VERSION
+        )
+        # If metadata is None covert to {}
+        metadata = metadata if metadata else {}
+        self._nested_update_metadata(self._gca_resource, metadata)
+        self._gca_resource.description = description
+
+    # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
+    def _init_with_resource_name(
+        self,
+        *,
+        context_name: str,
+    ):
+        """Initializes the Artifact instance using an existing resource.
+        Args:
+            context_name (str):
+                Context name with the following format, this is globally unique in a metadataStore:
+                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/contexts/<resource_id>.
+        """
+        super(BaseContextSchema, self).__init__(resource_name=context_name)
 
     def create(
         self,
@@ -97,15 +122,24 @@ def create(
             Context: Instantiated representation of the managed Metadata Context.
 
         """
-        return context.Context.create(
+        # Check if metadata exists to avoid proto read error
+        metadata = None
+        if self._gca_resource.metadata:
+            metadata = self.metadata
+
+        new_context = context.Context.create(
             resource_id=self.context_id,
             schema_title=self.schema_title,
             display_name=self.display_name,
             schema_version=self.schema_version,
             description=self.description,
-            metadata=self.metadata,
+            metadata=metadata,
             metadata_store_id=metadata_store_id,
             project=project,
             location=location,
             credentials=credentials,
         )
+
+        # Reinstantiate this class using the newly created resource.
+        self._init_with_resource_name(context_name=new_context.resource_name)
+        return self