Skip to content

Commit 3d08481

Browse files
ErnestoLomaErnestoLoma
committed
Updated Schema.create method to add schema to Database schemas set and updated Table.create method to add table to Schema tables set.
1 parent 956071d commit 3d08481

4 files changed

Lines changed: 94 additions & 98 deletions

File tree

‎pyatlan/generator/templates/entity.jinja2‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -447,10 +447,14 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
447447
{%- elif entity_def.name == "Schema" %}
448448
@classmethod
449449
# @validate_arguments()
450-
def create(cls, *, name: str, database_qualified_name: str)->{{ entity_def.name }}.Attributes:
450+
def create(
451+
cls, *, name: str, database_qualified_name: str
452+
) -> Schema.Attributes:
451453
if not name:
452454
raise ValueError("name cannot be blank")
453-
validate_required_fields(["database_qualified_name"], [database_qualified_name])
455+
validate_required_fields(
456+
["database_qualified_name"], [database_qualified_name]
457+
)
454458
fields = database_qualified_name.split("/")
455459
if len(fields) != 4:
456460
raise ValueError("Invalid database_qualified_name")
@@ -465,11 +469,12 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
465469
database_qualified_name=database_qualified_name,
466470
qualified_name=f"{database_qualified_name}/{name}",
467471
connector_name=connector_type.value,
472+
database=Database.ref_by_qualified_name(database_qualified_name),
468473
)
469474
{%- elif entity_def.name == "Table" or entity_def.name == "View" %}
470475
@classmethod
471476
# @validate_arguments()
472-
def create(cls, *, name: str, schema_qualified_name: str)->{{ entity_def.name }}.Attributes:
477+
def create(cls, *, name: str, schema_qualified_name: str) -> {{ entity_def.name }}.Attributes:
473478
if not name:
474479
raise ValueError("name cannot be blank")
475480
validate_required_fields(["schema_qualified_name"], [schema_qualified_name])
@@ -489,6 +494,7 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
489494
schema_qualified_name=schema_qualified_name,
490495
schema_name=fields[4],
491496
connector_name=connector_type.value,
497+
atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name),
492498
)
493499
{%- elif entity_def.name == "S3Bucket" %}
494500
@classmethod

‎pyatlan/model/assets.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9098,6 +9098,7 @@ def create(cls, *, name: str, schema_qualified_name: str) -> Table.Attributes:
90989098
schema_qualified_name=schema_qualified_name,
90999099
schema_name=fields[4],
91009100
connector_name=connector_type.value,
9101+
atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name),
91019102
)
91029103

91039104
attributes: "Table.Attributes" = Field(
@@ -10148,6 +10149,7 @@ def create(
1014810149
database_qualified_name=database_qualified_name,
1014910150
qualified_name=f"{database_qualified_name}/{name}",
1015010151
connector_name=connector_type.value,
10152+
database=Database.ref_by_qualified_name(database_qualified_name),
1015110153
)
1015210154

1015310155
attributes: "Schema.Attributes" = Field(
@@ -10785,6 +10787,7 @@ def create(cls, *, name: str, schema_qualified_name: str) -> View.Attributes:
1078510787
schema_qualified_name=schema_qualified_name,
1078610788
schema_name=fields[4],
1078710789
connector_name=connector_type.value,
10790+
atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name),
1078810791
)
1078910792

1079010793
attributes: "View.Attributes" = Field(

‎pyatlan/model/enums.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,17 @@ class AtlanConnectionCategory(Enum):
198198
class AtlanConnectorType(str, Enum):
199199
category: AtlanConnectionCategory
200200

201+
@classmethod
202+
def _get_connector_type_from_qualified_name(
203+
cls, qualified_name: str
204+
) -> "AtlanConnectorType":
205+
tokens = qualified_name.split("/")
206+
if len(tokens) > 1:
207+
return AtlanConnectorType[tokens[1].upper()]
208+
raise ValueError(
209+
f"Could not determine AtlanConnectorType from {qualified_name}"
210+
)
211+
201212
def __new__(
202213
cls, value: str, category: AtlanConnectionCategory
203214
) -> "AtlanConnectorType":

‎tests/integration/test_entity_model.py‎

Lines changed: 71 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import os
44
import random
55
import string
6-
import time
76

87
import pytest
98
import requests
@@ -42,6 +41,7 @@
4241
"4af8d57c-61ef-4b57-983c-eff20e6d08b5",
4342
"57f5463d-cc2a-4859-bf28-e4fa52002e8e",
4443
}
44+
TEMP_CONNECTION_GUID = "b3a5c49a-0c7c-4e66-8453-f4da8d9ce222"
4545

4646

4747
@pytest.fixture(scope="module")
@@ -85,13 +85,13 @@ def get_environment_variable(name) -> str:
8585
return ret_value
8686

8787

88-
@pytest.fixture()
88+
@pytest.fixture(scope="session")
8989
def increment_counter():
90-
i = random.randint(0, 1000)
90+
i = 700
9191

9292
def increment():
9393
nonlocal i
94-
i += 1
94+
i += 20
9595
return i
9696

9797
return increment
@@ -523,23 +523,23 @@ def test_create_connection(client: AtlanClient, increment_counter):
523523
assert c.guid == guid
524524

525525

526-
@pytest.mark.skip("Connection creation is still intermittently failing")
527526
def test_create_database(client: AtlanClient, increment_counter):
528527
role = RoleCache.get_id_for_name("$admin")
529528
assert role
530529
suffix = increment_counter()
531-
connection = Connection.create(
532-
name=f"Integration {suffix}",
533-
connector_type=AtlanConnectorType.SNOWFLAKE,
534-
admin_roles=[role],
535-
admin_groups=["admin"],
536-
)
537-
response = client.upsert(connection)
538-
assert response.mutated_entities
539-
assert response.mutated_entities.CREATE
540-
assert isinstance(response.mutated_entities.CREATE[0], Connection)
541-
connection = response.mutated_entities.CREATE[0]
542-
connection = client.get_asset_by_guid(connection.guid, Connection)
530+
# connection = Connection.create(
531+
# name=f"Integration {suffix}",
532+
# connector_type=AtlanConnectorType.SNOWFLAKE,
533+
# admin_roles=[role],
534+
# admin_groups=["admin"],
535+
# )
536+
# response = client.upsert(connection)
537+
# assert response.mutated_entities
538+
# assert response.mutated_entities.CREATE
539+
# assert isinstance(response.mutated_entities.CREATE[0], Connection)
540+
# connection = response.mutated_entities.CREATE[0]
541+
# connection = client.get_asset_by_guid(connection.guid, Connection)
542+
connection = client.get_asset_by_guid(TEMP_CONNECTION_GUID, Connection)
543543
database = Database.create(
544544
name=f"Integration_{suffix}",
545545
connection_qualified_name=connection.attributes.qualified_name,
@@ -550,122 +550,98 @@ def test_create_database(client: AtlanClient, increment_counter):
550550
assert len(response.mutated_entities.CREATE) == 1
551551
assert isinstance(response.mutated_entities.CREATE[0], Database)
552552
assert response.guid_assignments
553-
assert database.guid in response.guid_assignments
554-
guid = response.guid_assignments[database.guid]
555553
database = response.mutated_entities.CREATE[0]
556-
assert guid == database.guid
557-
database = client.get_asset_by_guid(guid, Database)
558-
assert isinstance(database, Database)
559-
assert guid == database.guid
554+
client.get_asset_by_guid(database.guid, Database)
560555

561556

562-
@pytest.mark.skip("Connection creation is still intermittently failing")
563557
def test_create_schema(client: AtlanClient, increment_counter):
564558
role = RoleCache.get_id_for_name("$admin")
565559
assert role
566560
suffix = increment_counter()
567-
connection = Connection.create(
568-
name=f"Integration {suffix}",
569-
connector_type=AtlanConnectorType.SNOWFLAKE,
570-
admin_roles=[role],
571-
admin_groups=["admin"],
572-
)
573-
response = client.upsert(connection)
574-
assert response.mutated_entities
575-
assert response.mutated_entities.CREATE
576-
assert isinstance(response.mutated_entities.CREATE[0], Connection)
577-
connection = response.mutated_entities.CREATE[0]
578-
time.sleep(30)
579-
connection = client.get_asset_by_guid(connection.guid, Connection)
561+
# connection = Connection.create(
562+
# name=f"Integration {suffix}",
563+
# connector_type=AtlanConnectorType.SNOWFLAKE,
564+
# admin_roles=[role],
565+
# admin_groups=["admin"],
566+
# )
567+
# response = client.upsert(connection)
568+
# assert response.mutated_entities
569+
# assert response.mutated_entities.CREATE
570+
# assert isinstance(response.mutated_entities.CREATE[0], Connection)
571+
# connection = response.mutated_entities.CREATE[0]
572+
# time.sleep(30)
573+
connection = client.get_asset_by_guid(TEMP_CONNECTION_GUID, Connection)
580574
database = Database.create(
581575
name=f"Integration_{suffix}",
582576
connection_qualified_name=connection.attributes.qualified_name,
583577
)
584578
response = client.upsert(database)
585-
assert response.mutated_entities
586-
assert response.mutated_entities.CREATE
587-
assert isinstance(response.mutated_entities.CREATE[0], Database)
588-
database = response.mutated_entities.CREATE[0]
589-
time.sleep(3)
590-
database = client.get_asset_by_guid(database.guid, Database)
579+
assert (databases := response.assets_created(asset_type=Database))
580+
assert len(databases) == 1
581+
database = client.get_asset_by_guid(databases[0].guid, Database)
591582
schema = Schema.create(
592583
name=f"Integration_{suffix}",
593584
database_qualified_name=database.attributes.qualified_name,
594585
)
595586
response = client.upsert(schema)
596-
assert response.mutated_entities
597-
assert response.mutated_entities.CREATE
598-
assert len(response.mutated_entities.CREATE) == 1
599-
assert isinstance(response.mutated_entities.CREATE[0], Schema)
600-
assert response.guid_assignments
601-
assert schema.guid in response.guid_assignments
602-
guid = response.guid_assignments[schema.guid]
603-
schema = response.mutated_entities.CREATE[0]
604-
assert guid == schema.guid
605-
time.sleep(3)
606-
schema = client.get_asset_by_guid(guid, Schema)
607-
assert isinstance(schema, Schema)
608-
assert guid == schema.guid
587+
assert (schemas := response.assets_created(asset_type=Schema))
588+
assert len(schemas) == 1
589+
schema = client.get_asset_by_guid(schemas[0].guid, Schema)
590+
assert (databases := response.assets_updated(asset_type=Database))
591+
assert len(databases) == 1
592+
database = client.get_asset_by_guid(databases[0].guid, Database)
593+
assert database.attributes.schemas
594+
schemas = database.attributes.schemas
595+
assert len(schemas) == 1
596+
assert schemas[0].guid == schema.guid
609597

610598

611-
@pytest.mark.skip("Connection creation is still intermittently failing")
612599
def test_create_table(client: AtlanClient, increment_counter):
613600
role = RoleCache.get_id_for_name("$admin")
614601
assert role
615602
suffix = increment_counter()
616-
connection = Connection.create(
617-
name=f"Integration {suffix}",
618-
connector_type=AtlanConnectorType.SNOWFLAKE,
619-
admin_roles=[role],
620-
admin_groups=["admin"],
621-
)
622-
response = client.upsert(connection)
623-
assert response.mutated_entities
624-
assert response.mutated_entities.CREATE
625-
assert isinstance(response.mutated_entities.CREATE[0], Connection)
626-
connection = response.mutated_entities.CREATE[0]
627-
time.sleep(30)
628-
connection = client.get_asset_by_guid(connection.guid, Connection)
603+
# connection = Connection.create(
604+
# name=f"Integration {suffix}",
605+
# connector_type=AtlanConnectorType.SNOWFLAKE,
606+
# admin_roles=[role],
607+
# admin_groups=["admin"],
608+
# )
609+
# response = client.upsert(connection)
610+
# assert response.mutated_entities
611+
# assert response.mutated_entities.CREATE
612+
# assert isinstance(response.mutated_entities.CREATE[0], Connection)
613+
# connection = response.mutated_entities.CREATE[0]
614+
# time.sleep(30)
615+
connection = client.get_asset_by_guid(TEMP_CONNECTION_GUID, Connection)
629616
database = Database.create(
630617
name=f"Integration_{suffix}",
631618
connection_qualified_name=connection.attributes.qualified_name,
632619
)
633620
response = client.upsert(database)
634-
assert response.mutated_entities
635-
assert response.mutated_entities.CREATE
636-
assert isinstance(response.mutated_entities.CREATE[0], Database)
637-
database = response.mutated_entities.CREATE[0]
638-
time.sleep(3)
639-
database = client.get_asset_by_guid(database.guid, Database)
621+
assert (databases := response.assets_created(asset_type=Database))
622+
database = client.get_asset_by_guid(databases[0].guid, Database)
640623
schema = Schema.create(
641624
name=f"Integration_{suffix}",
642625
database_qualified_name=database.attributes.qualified_name,
643626
)
644627
response = client.upsert(schema)
645-
assert response.mutated_entities
646-
assert response.mutated_entities.CREATE
647-
assert isinstance(response.mutated_entities.CREATE[0], Schema)
648-
schema = response.mutated_entities.CREATE[0]
649-
time.sleep(3)
650-
schema = client.get_asset_by_guid(schema.guid, Schema)
628+
assert (schemas := response.assets_created(asset_type=Schema))
629+
schema = client.get_asset_by_guid(schemas[0].guid, Schema)
651630
table = Table.create(
652631
name=f"Integration_{suffix}",
653632
schema_qualified_name=schema.attributes.qualified_name,
654633
)
655634
response = client.upsert(table)
656-
assert response.mutated_entities
657-
assert response.mutated_entities.CREATE
658-
assert len(response.mutated_entities.CREATE) == 1
659-
assert isinstance(response.mutated_entities.CREATE[0], Table)
660-
assert response.guid_assignments
661-
assert table.guid in response.guid_assignments
662-
guid = response.guid_assignments[table.guid]
663-
table = response.mutated_entities.CREATE[0]
664-
assert guid == table.guid
665-
time.sleep(3)
666-
table = client.get_asset_by_guid(guid, Table)
667-
assert isinstance(table, Table)
668-
assert guid == table.guid
635+
assert (tables := response.assets_created(asset_type=Table))
636+
assert len(tables) == 1
637+
table = client.get_asset_by_guid(guid=tables[0].guid, asset_type=Table)
638+
assert (schemas := response.assets_updated(asset_type=Schema))
639+
assert len(schemas) == 1
640+
schema = client.get_asset_by_guid(guid=schemas[0].guid, asset_type=Schema)
641+
assert schema.attributes.tables
642+
tables = schema.attributes.tables
643+
assert len(tables) == 1
644+
assert tables[0].guid == table.guid
669645

670646

671647
def test_get_by_qualified_name(client: AtlanClient):

0 commit comments

Comments
 (0)