Add create method to Column

ErnestoLoma · ErnestoLoma · commit 187d02907e14 · 2023-04-12T19:38:38.000+03:00
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,3 +1,9 @@
+## 0.0.22 (Apr 12, 2023)
+
+* Added create metho 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
@@ -496,6 +496,63 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 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
@@ -10022,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 "
diff --git a/pyatlan/version.txt b/pyatlan/version.txt
@@ -1 +1 @@
-0.0.21
+0.0.22
diff --git a/tests/integration/test_entity_model.py b/tests/integration/test_entity_model.py
@@ -187,6 +187,7 @@ def cleanup(atlan_host, headers, atlan_api_key):
         "Database",
         "Connection",
         "View",
+        "Column",
     ]
     for type_name in type_names:
         print()
@@ -670,6 +671,61 @@ def test_create_view(client: AtlanClient, increment_counter):
     assert guid == view.guid
 
 
+def test_create_column(client: AtlanClient, increment_counter):
+    role = RoleCache.get_id_for_name("$admin")
+    assert role
+    suffix = increment_counter()
+    # connection = Connection.create(
+    #     name=f"Integration {suffix}",
+    #     connector_type=AtlanConnectorType.SNOWFLAKE,
+    #     admin_roles=[role],
+    #     admin_groups=["admin"],
+    # )
+    # response = client.upsert(connection)
+    # assert response.mutated_entities
+    # assert response.mutated_entities.CREATE
+    # assert isinstance(response.mutated_entities.CREATE[0], Connection)
+    # connection = response.mutated_entities.CREATE[0]
+    # time.sleep(30)
+    connection = client.get_asset_by_guid(TEMP_CONNECTION_GUID, Connection)
+    database = Database.create(
+        name=f"Integration_{suffix}",
+        connection_qualified_name=connection.attributes.qualified_name,
+    )
+    response = client.upsert(database)
+    assert (databases := response.assets_created(asset_type=Database))
+    database = client.get_asset_by_guid(databases[0].guid, Database)
+    schema = Schema.create(
+        name=f"Integration_{suffix}",
+        database_qualified_name=database.attributes.qualified_name,
+    )
+    response = client.upsert(schema)
+    assert (schemas := response.assets_created(asset_type=Schema))
+    schema = client.get_asset_by_guid(schemas[0].guid, Schema)
+    table = Table.create(
+        name=f"Integration_{suffix}",
+        schema_qualified_name=schema.attributes.qualified_name,
+    )
+    response = client.upsert(table)
+    assert (tables := response.assets_created(asset_type=Table))
+    table = client.get_asset_by_guid(guid=tables[0].guid, asset_type=Table)
+    column = Column.create(
+        name=f"Integration_{suffix}_column",
+        parent_qualified_name=table.qualified_name,
+        parent_type=Table,
+        order=1,
+    )
+    response = client.upsert(column)
+    assert (columns := response.assets_created(asset_type=Column))
+    assert len(columns) == 1
+    column = client.get_asset_by_guid(asset_type=Column, guid=columns[0].guid)
+    table = client.get_asset_by_guid(asset_type=Table, guid=table.guid)
+    assert table.attributes.columns
+    columns = table.attributes.columns
+    assert len(columns) == 1
+    assert columns[0].guid == column.guid
+
+
 def test_add_and_remove_classifications(client: AtlanClient):
     glossary = AtlasGlossary.create(name="Integration Classification Test")
     glossary.attributes.user_description = "This is a description of the glossary"
