Merge pull request #38 from atlanhq/update_create_functions

Update create functions
atlanhq · Apr 12, 2023 · f363382 · f363382
2 parents 956071d + 4038555
commit f363382
Show file tree

Hide file tree

Showing 6 changed files with 271 additions and 99 deletions.
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,3 +1,9 @@
+## 0.0.22 (Apr 12, 2023)
+
+* Added create method to Column
+* Updated Schema.create method to add schema to Database schemas collection
+* Updated Table.create method to add table to Schema tables collection
+
 ## 0.0.21 (Apr 11, 2023)
 
 * Added relation_attributes parameter to IndexSearchRequest to specify the attributes to be included in each relationship that is included in the results of the search

diff --git a/pyatlan/generator/templates/entity.jinja2 b/pyatlan/generator/templates/entity.jinja2
@@ -447,10 +447,14 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
         {%-  elif entity_def.name == "Schema" %}
         @classmethod
         # @validate_arguments()
-        def create(cls, *, name: str, database_qualified_name: str)->{{ entity_def.name }}.Attributes:
+        def create(
+            cls, *, name: str, database_qualified_name: str
+        ) -> Schema.Attributes:
             if not name:
                 raise ValueError("name cannot be blank")
-            validate_required_fields(["database_qualified_name"], [database_qualified_name])
+            validate_required_fields(
+                ["database_qualified_name"], [database_qualified_name]
+            )
             fields = database_qualified_name.split("/")
             if len(fields) != 4:
                 raise ValueError("Invalid database_qualified_name")
@@ -465,11 +469,12 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 database_qualified_name=database_qualified_name,
                 qualified_name=f"{database_qualified_name}/{name}",
                 connector_name=connector_type.value,
+                database=Database.ref_by_qualified_name(database_qualified_name),
             )
         {%- elif entity_def.name == "Table" or entity_def.name == "View" %}
         @classmethod
         # @validate_arguments()
-        def create(cls, *, name: str, schema_qualified_name: str)->{{ entity_def.name }}.Attributes:
+        def create(cls, *, name: str, schema_qualified_name: str) -> {{ entity_def.name }}.Attributes:
             if not name:
                 raise ValueError("name cannot be blank")
             validate_required_fields(["schema_qualified_name"], [schema_qualified_name])
@@ -489,7 +494,65 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 schema_qualified_name=schema_qualified_name,
                 schema_name=fields[4],
                 connector_name=connector_type.value,
+                atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name),
             )
+        {%- elif entity_def.name == "Column" %}
+        @classmethod
+        # @validate_arguments()
+        def create(
+            cls, *, name: str, parent_qualified_name: str, parent_type: type, order: int
+        ) -> Column.Attributes:
+            if not name:
+                raise ValueError("name cannot be blank")
+            validate_required_fields(["parent_qualified_name"], [parent_qualified_name])
+            fields = parent_qualified_name.split("/")
+            if len(fields) != 6:
+                raise ValueError("Invalid parent_qualified_name")
+            try:
+                connector_type = AtlanConnectorType(fields[1])  # type:ignore
+            except ValueError as e:
+                raise ValueError("Invalid parent_qualified_name") from e
+            ret_value = Column.Attributes(
+                name=name,
+                qualified_name=f"{parent_qualified_name}/{name}",
+                connector_name=connector_type.value,
+                schema_name=fields[4],
+                schema_qualified_name=f"{fields[0]}/{fields[1]}/{fields[2]}/{fields[3]}/{fields[4]}",
+                database_name=fields[3],
+                database_qualified_name=f"{fields[0]}/{fields[1]}/{fields[2]}/{fields[3]}",
+                connection_qualified_name=f"{fields[0]}/{fields[1]}/{fields[2]}",
+                order=order,
+            )
+            if parent_type == Table:
+                ret_value.table_qualified_name = parent_qualified_name
+                ret_value.table = Table.ref_by_qualified_name(parent_qualified_name)
+            elif parent_type == View:
+                ret_value.view_qualified_name = parent_qualified_name
+                ret_value.view = View.ref_by_qualified_name(parent_qualified_name)
+            elif parent_type == MaterialisedView:
+                ret_value.view_qualified_name = parent_qualified_name
+                ret_value.materialised_view = MaterialisedView.ref_by_qualified_name(
+                    parent_qualified_name
+                )
+            else:
+                raise ValueError(
+                    f"parent_type must be either Table, View or MaterializeView"
+                )
+            return ret_value
+
+    @classmethod
+    # @validate_arguments()
+    def create(
+        cls, *, name: str, parent_qualified_name: str, parent_type: type, order: int
+    ) -> Column:
+        return Column(
+            attributes=Column.Attributes.create(
+                name=name,
+                parent_qualified_name=parent_qualified_name,
+                parent_type=parent_type,
+                order=order,
+            )
+        )
         {%- elif entity_def.name == "S3Bucket" %}
         @classmethod
         # @validate_arguments()

diff --git a/pyatlan/model/assets.py b/pyatlan/model/assets.py
@@ -9098,6 +9098,7 @@ def create(cls, *, name: str, schema_qualified_name: str) -> Table.Attributes:
                 schema_qualified_name=schema_qualified_name,
                 schema_name=fields[4],
                 connector_name=connector_type.value,
+                atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name),
             )
 
     attributes: "Table.Attributes" = Field(
@@ -10021,6 +10022,63 @@ class Attributes(SQL.Attributes):
             None, description="", alias="columnDbtModelColumns"
         )  # relationship
 
+        @classmethod
+        # @validate_arguments()
+        def create(
+            cls, *, name: str, parent_qualified_name: str, parent_type: type, order: int
+        ) -> Column.Attributes:
+            if not name:
+                raise ValueError("name cannot be blank")
+            validate_required_fields(["parent_qualified_name"], [parent_qualified_name])
+            fields = parent_qualified_name.split("/")
+            if len(fields) != 6:
+                raise ValueError("Invalid parent_qualified_name")
+            try:
+                connector_type = AtlanConnectorType(fields[1])  # type:ignore
+            except ValueError as e:
+                raise ValueError("Invalid parent_qualified_name") from e
+            ret_value = Column.Attributes(
+                name=name,
+                qualified_name=f"{parent_qualified_name}/{name}",
+                connector_name=connector_type.value,
+                schema_name=fields[4],
+                schema_qualified_name=f"{fields[0]}/{fields[1]}/{fields[2]}/{fields[3]}/{fields[4]}",
+                database_name=fields[3],
+                database_qualified_name=f"{fields[0]}/{fields[1]}/{fields[2]}/{fields[3]}",
+                connection_qualified_name=f"{fields[0]}/{fields[1]}/{fields[2]}",
+                order=order,
+            )
+            if parent_type == Table:
+                ret_value.table_qualified_name = parent_qualified_name
+                ret_value.table = Table.ref_by_qualified_name(parent_qualified_name)
+            elif parent_type == View:
+                ret_value.view_qualified_name = parent_qualified_name
+                ret_value.view = View.ref_by_qualified_name(parent_qualified_name)
+            elif parent_type == MaterialisedView:
+                ret_value.view_qualified_name = parent_qualified_name
+                ret_value.materialised_view = MaterialisedView.ref_by_qualified_name(
+                    parent_qualified_name
+                )
+            else:
+                raise ValueError(
+                    "parent_type must be either Table, View or MaterializeView"
+                )
+            return ret_value
+
+    @classmethod
+    # @validate_arguments()
+    def create(
+        cls, *, name: str, parent_qualified_name: str, parent_type: type, order: int
+    ) -> Column:
+        return Column(
+            attributes=Column.Attributes.create(
+                name=name,
+                parent_qualified_name=parent_qualified_name,
+                parent_type=parent_type,
+                order=order,
+            )
+        )
+
     attributes: "Column.Attributes" = Field(
         None,
         description="Map of attributes in the instance and their values. The specific keys of this map will vary by "
@@ -10148,6 +10206,7 @@ def create(
                 database_qualified_name=database_qualified_name,
                 qualified_name=f"{database_qualified_name}/{name}",
                 connector_name=connector_type.value,
+                database=Database.ref_by_qualified_name(database_qualified_name),
             )
 
     attributes: "Schema.Attributes" = Field(
@@ -10785,6 +10844,7 @@ def create(cls, *, name: str, schema_qualified_name: str) -> View.Attributes:
                 schema_qualified_name=schema_qualified_name,
                 schema_name=fields[4],
                 connector_name=connector_type.value,
+                atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name),
             )
 
     attributes: "View.Attributes" = Field(

diff --git a/pyatlan/model/enums.py b/pyatlan/model/enums.py
@@ -198,6 +198,17 @@ class AtlanConnectionCategory(Enum):
 class AtlanConnectorType(str, Enum):
     category: AtlanConnectionCategory
 
+    @classmethod
+    def _get_connector_type_from_qualified_name(
+        cls, qualified_name: str
+    ) -> "AtlanConnectorType":
+        tokens = qualified_name.split("/")
+        if len(tokens) > 1:
+            return AtlanConnectorType[tokens[1].upper()]
+        raise ValueError(
+            f"Could not determine AtlanConnectorType from {qualified_name}"
+        )
+
     def __new__(
         cls, value: str, category: AtlanConnectionCategory
     ) -> "AtlanConnectorType":

diff --git a/pyatlan/version.txt b/pyatlan/version.txt
@@ -1 +1 @@
-0.0.21
+0.0.22