From be516ce35613789a2f16e75185fa1d6de1a9fcc7 Mon Sep 17 00:00:00 2001
From: rupali-atlan <rupali.kishore@atlan.com>
Date: Fri, 10 Jan 2025 19:02:50 +0530
Subject: [PATCH 1/6] added custom creator methods,generator templates, unit
and integration test
---
.../methods/asset/custom_entity.jinja2 | 11 ++
.../methods/attribute/custom_entity.jinja2 | 17 ++
pyatlan/model/assets/__init__.py | 8 +-
pyatlan/model/assets/__init__.pyi | 8 +-
pyatlan/model/assets/custom.py | 33 ++++
pyatlan/model/assets/custom_entity.py | 172 ++++++++++++++++++
pyatlan/model/constants.py | 1 +
pyatlan/model/enums.py | 2 +
pyatlan/model/typedef.py | 1 +
tests/integration/custom_asset_test.py | 96 ++++++++++
tests/unit/model/constants.py | 6 +
tests/unit/model/custom_entity_test.py | 69 +++++++
tests/unit/test_model.py | 2 +
13 files changed, 422 insertions(+), 4 deletions(-)
create mode 100644 pyatlan/generator/templates/methods/asset/custom_entity.jinja2
create mode 100644 pyatlan/generator/templates/methods/attribute/custom_entity.jinja2
create mode 100644 pyatlan/model/assets/custom.py
create mode 100644 pyatlan/model/assets/custom_entity.py
create mode 100644 tests/integration/custom_asset_test.py
create mode 100644 tests/unit/model/custom_entity_test.py
diff --git a/pyatlan/generator/templates/methods/asset/custom_entity.jinja2 b/pyatlan/generator/templates/methods/asset/custom_entity.jinja2
new file mode 100644
index 000000000..f78fb35d8
--- /dev/null
+++ b/pyatlan/generator/templates/methods/asset/custom_entity.jinja2
@@ -0,0 +1,11 @@
+
+ @classmethod
+ @init_guid
+ def creator(cls, *, name: str, connection_qualified_name: str) -> CustomEntity:
+ validate_required_fields(
+ ["name", "connection_qualified_name"], [name, connection_qualified_name]
+ )
+ attributes = CustomEntity.Attributes.create(
+ name=name, connection_qualified_name=connection_qualified_name
+ )
+ return cls(attributes=attributes)
diff --git a/pyatlan/generator/templates/methods/attribute/custom_entity.jinja2 b/pyatlan/generator/templates/methods/attribute/custom_entity.jinja2
new file mode 100644
index 000000000..be2a575e5
--- /dev/null
+++ b/pyatlan/generator/templates/methods/attribute/custom_entity.jinja2
@@ -0,0 +1,17 @@
+
+ @classmethod
+ @init_guid
+ def create(
+ cls, *, name: str, connection_qualified_name: str
+ ) -> CustomEntity.Attributes:
+ validate_required_fields(
+ ["name", "connection_qualified_name"], [name, connection_qualified_name]
+ )
+ return CustomEntity.Attributes(
+ name=name,
+ qualified_name=f"{connection_qualified_name}/{name}",
+ connection_qualified_name=connection_qualified_name,
+ connector_name=AtlanConnectorType.get_connector_name(
+ connection_qualified_name
+ ),
+ )
diff --git a/pyatlan/model/assets/__init__.py b/pyatlan/model/assets/__init__.py
index 54d13cfe2..f87a62701 100644
--- a/pyatlan/model/assets/__init__.py
+++ b/pyatlan/model/assets/__init__.py
@@ -64,10 +64,10 @@
"Function",
"TablePartition",
"Column",
- "DatabricksUnityCatalogTag",
"SnowflakeStream",
- "CalculationView",
+ "DatabricksUnityCatalogTag",
"Database",
+ "CalculationView",
"Procedure",
"SnowflakeTag",
"CosmosMongoDB",
@@ -111,6 +111,8 @@
"CosmosMongoDBAccount",
"CosmosMongoDBCollection",
"CosmosMongoDBDatabase",
+ "Custom",
+ "CustomEntity",
],
"task": ["Task"],
"data_set": ["DataSet"],
@@ -133,6 +135,7 @@
"object_store": ["ObjectStore"],
"saa_s": ["SaaS"],
"multi_dimensional_dataset": ["MultiDimensionalDataset"],
+ "custom": ["Custom"],
"event_store": ["EventStore"],
"insight": ["Insight"],
"a_p_i": ["API"],
@@ -169,6 +172,7 @@
"cube_hierarchy": ["CubeHierarchy"],
"cube_field": ["CubeField"],
"cube_dimension": ["CubeDimension"],
+ "custom_entity": ["CustomEntity"],
"bigquery_tag": ["BigqueryTag"],
"kafka": ["Kafka"],
"azure_service_bus": ["AzureServiceBus"],
diff --git a/pyatlan/model/assets/__init__.pyi b/pyatlan/model/assets/__init__.pyi
index 686622375..6653f1fdf 100644
--- a/pyatlan/model/assets/__init__.pyi
+++ b/pyatlan/model/assets/__init__.pyi
@@ -61,10 +61,10 @@ __all__ = [
"Function",
"TablePartition",
"Column",
- "DatabricksUnityCatalogTag",
"SnowflakeStream",
- "CalculationView",
+ "DatabricksUnityCatalogTag",
"Database",
+ "CalculationView",
"Procedure",
"SnowflakeTag",
"CosmosMongoDB",
@@ -129,6 +129,7 @@ __all__ = [
"ObjectStore",
"SaaS",
"MultiDimensionalDataset",
+ "Custom",
"EventStore",
"NoSQL",
"Insight",
@@ -166,6 +167,7 @@ __all__ = [
"CubeHierarchy",
"CubeField",
"CubeDimension",
+ "CustomEntity",
"BigqueryTag",
"Kafka",
"AzureServiceBus",
@@ -474,6 +476,8 @@ from .cube import Cube
from .cube_dimension import CubeDimension
from .cube_field import CubeField
from .cube_hierarchy import CubeHierarchy
+from .custom import Custom
+from .custom_entity import CustomEntity
from .data_set import DataSet
from .data_studio import DataStudio
from .data_studio_asset import DataStudioAsset
diff --git a/pyatlan/model/assets/custom.py b/pyatlan/model/assets/custom.py
new file mode 100644
index 000000000..818215fef
--- /dev/null
+++ b/pyatlan/model/assets/custom.py
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 Atlan Pte. Ltd.
+
+
+from __future__ import annotations
+
+from typing import ClassVar, List
+
+from pydantic.v1 import Field, validator
+
+from .core.catalog import Catalog
+
+
+class Custom(Catalog):
+ """Description"""
+
+ type_name: str = Field(default="Custom", allow_mutation=False)
+
+ @validator("type_name")
+ def validate_type_name(cls, v):
+ if v != "Custom":
+ raise ValueError("must be Custom")
+ return v
+
+ def __setattr__(self, name, value):
+ if name in Custom._convenience_properties:
+ return object.__setattr__(self, name, value)
+ super().__setattr__(name, value)
+
+ _convenience_properties: ClassVar[List[str]] = []
+
+
+Custom.Attributes.update_forward_refs()
diff --git a/pyatlan/model/assets/custom_entity.py b/pyatlan/model/assets/custom_entity.py
new file mode 100644
index 000000000..2630aa597
--- /dev/null
+++ b/pyatlan/model/assets/custom_entity.py
@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 Atlan Pte. Ltd.
+
+
+from __future__ import annotations
+
+from typing import ClassVar, List, Optional
+from pyatlan.model.enums import AtlanConnectorType
+from pydantic.v1 import Field, validator
+from pyatlan.utils import init_guid, validate_required_fields
+from pyatlan.model.fields.atlan_fields import RelationField
+
+from .custom import Custom
+
+
+class CustomEntity(Custom):
+ """Description"""
+
+ @classmethod
+ @init_guid
+ def creator(cls, *, name: str, connection_qualified_name: str) -> CustomEntity:
+ validate_required_fields(
+ ["name", "connection_qualified_name"], [name, connection_qualified_name]
+ )
+ attributes = CustomEntity.Attributes.create(
+ name=name, connection_qualified_name=connection_qualified_name
+ )
+ return cls(attributes=attributes)
+
+ type_name: str = Field(default="CustomEntity", allow_mutation=False)
+
+ @validator("type_name")
+ def validate_type_name(cls, v):
+ if v != "CustomEntity":
+ raise ValueError("must be CustomEntity")
+ return v
+
+ def __setattr__(self, name, value):
+ if name in CustomEntity._convenience_properties:
+ return object.__setattr__(self, name, value)
+ super().__setattr__(name, value)
+
+ CUSTOM_PARENT_ENTITY: ClassVar[RelationField] = RelationField("customParentEntity")
+ """
+ TBC
+ """
+ CUSTOM_CHILD_ENTITIES: ClassVar[RelationField] = RelationField(
+ "customChildEntities"
+ )
+ """
+ TBC
+ """
+ CUSTOM_RELATED_TO_ENTITIES: ClassVar[RelationField] = RelationField(
+ "customRelatedToEntities"
+ )
+ """
+ TBC
+ """
+ CUSTOM_RELATED_FROM_ENTITIES: ClassVar[RelationField] = RelationField(
+ "customRelatedFromEntities"
+ )
+ """
+ TBC
+ """
+
+ _convenience_properties: ClassVar[List[str]] = [
+ "custom_parent_entity",
+ "custom_child_entities",
+ "custom_related_to_entities",
+ "custom_related_from_entities",
+ ]
+
+ @property
+ def custom_parent_entity(self) -> Optional[CustomEntity]:
+ return None if self.attributes is None else self.attributes.custom_parent_entity
+
+ @custom_parent_entity.setter
+ def custom_parent_entity(self, custom_parent_entity: Optional[CustomEntity]):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.custom_parent_entity = custom_parent_entity
+
+ @property
+ def custom_child_entities(self) -> Optional[List[CustomEntity]]:
+ return (
+ None if self.attributes is None else self.attributes.custom_child_entities
+ )
+
+ @custom_child_entities.setter
+ def custom_child_entities(
+ self, custom_child_entities: Optional[List[CustomEntity]]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.custom_child_entities = custom_child_entities
+
+ @property
+ def custom_related_to_entities(self) -> Optional[List[CustomEntity]]:
+ return (
+ None
+ if self.attributes is None
+ else self.attributes.custom_related_to_entities
+ )
+
+ @custom_related_to_entities.setter
+ def custom_related_to_entities(
+ self, custom_related_to_entities: Optional[List[CustomEntity]]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.custom_related_to_entities = custom_related_to_entities
+
+ @property
+ def custom_related_from_entities(self) -> Optional[List[CustomEntity]]:
+ return (
+ None
+ if self.attributes is None
+ else self.attributes.custom_related_from_entities
+ )
+
+ @custom_related_from_entities.setter
+ def custom_related_from_entities(
+ self, custom_related_from_entities: Optional[List[CustomEntity]]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.custom_related_from_entities = custom_related_from_entities
+
+ class Attributes(Custom.Attributes):
+ custom_parent_entity: Optional[CustomEntity] = Field(
+ default=None, description=""
+ ) # relationship
+ custom_child_entities: Optional[List[CustomEntity]] = Field(
+ default=None, description=""
+ ) # relationship
+ custom_related_to_entities: Optional[List[CustomEntity]] = Field(
+ default=None, description=""
+ ) # relationship
+ custom_related_from_entities: Optional[List[CustomEntity]] = Field(
+ default=None, description=""
+ ) # relationship
+
+
+ @classmethod
+ @init_guid
+ def create(
+ cls, *, name: str, connection_qualified_name: str
+ ) -> CustomEntity.Attributes:
+ validate_required_fields(
+ ["name", "connection_qualified_name"], [name, connection_qualified_name]
+ )
+ return CustomEntity.Attributes(
+ name=name,
+ qualified_name=f"{connection_qualified_name}/{name}",
+ connection_qualified_name=connection_qualified_name,
+ connector_name=AtlanConnectorType.get_connector_name(
+ connection_qualified_name
+ ),
+ )
+
+
+ attributes: CustomEntity.Attributes = Field(
+ default_factory=lambda: CustomEntity.Attributes(),
+ description=(
+ "Map of attributes in the instance and their values. "
+ "The specific keys of this map will vary by type, "
+ "so are described in the sub-types of this schema."
+ ),
+ )
+
+
+CustomEntity.Attributes.update_forward_refs()
diff --git a/pyatlan/model/constants.py b/pyatlan/model/constants.py
index 34ec484f2..072acf21d 100644
--- a/pyatlan/model/constants.py
+++ b/pyatlan/model/constants.py
@@ -42,6 +42,7 @@
"Column",
"ColumnProcess",
"Connection",
+ "CustomEntity",
"DataStudioAsset",
"Database",
"DbtColumnProcess",
diff --git a/pyatlan/model/enums.py b/pyatlan/model/enums.py
index 0deeec6f1..16aaf2089 100644
--- a/pyatlan/model/enums.py
+++ b/pyatlan/model/enums.py
@@ -137,6 +137,7 @@ class AtlanConnectionCategory(str, Enum):
DATA_QUALITY = "data-quality"
SCHEMA_REGISTRY = "schema-registry"
APP = "app"
+ CUSTOM = "custom"
class AtlanConnectorType(str, Enum):
@@ -342,6 +343,7 @@ def get_connector_name(
AWS_ECS = ("aws-ecs", AtlanConnectionCategory.ELT)
AWS_LAMBDA = ("aws-lambda", AtlanConnectionCategory.ELT)
AWS_SAGEMAKER = ("aws-sagemaker", AtlanConnectionCategory.ELT)
+ CUSTOM = ("custom", AtlanConnectionCategory.CUSTOM)
class AtlanCustomAttributePrimitiveType(str, Enum):
diff --git a/pyatlan/model/typedef.py b/pyatlan/model/typedef.py
index dd19c92bb..c3720101a 100644
--- a/pyatlan/model/typedef.py
+++ b/pyatlan/model/typedef.py
@@ -52,6 +52,7 @@
"Column",
"ColumnProcess",
"Connection",
+ "CustomEntity",
"DataStudioAsset",
"Database",
"DbtColumnProcess",
diff --git a/tests/integration/custom_asset_test.py b/tests/integration/custom_asset_test.py
new file mode 100644
index 000000000..fa84bcc91
--- /dev/null
+++ b/tests/integration/custom_asset_test.py
@@ -0,0 +1,96 @@
+from typing import Generator
+
+import pytest
+
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.model.assets import (
+ Connection, CustomEntity
+)
+from pyatlan.model.enums import (
+ AtlanConnectorType,
+ EntityStatus,
+)
+from pyatlan.model.response import AssetMutationResponse
+from tests.integration.client import TestId, delete_asset
+from tests.integration.connection_test import create_connection
+from tests.integration.utils import block
+
+MODULE_NAME = TestId.make_unique("CUSTOM")
+
+CONNECTOR_TYPE = AtlanConnectorType.CUSTOM
+CUSTOM_ENTITY_NAME = f"{MODULE_NAME}-custom-entity"
+
+response = block(AtlanClient(), AssetMutationResponse())
+
+
+@pytest.fixture(scope="module")
+def connection(client: AtlanClient) -> Generator[Connection, None, None]:
+ result = create_connection(
+ client=client, name=MODULE_NAME, connector_type=CONNECTOR_TYPE
+ )
+ yield result
+ delete_asset(client, guid=result.guid, asset_type=Connection)
+
+
+@pytest.fixture(scope="module")
+def custom_entity(
+ client: AtlanClient, connection: Connection
+) -> Generator[CustomEntity, None, None]:
+ assert connection.qualified_name
+ to_create = CustomEntity.creator(
+ name=CUSTOM_ENTITY_NAME, connection_qualified_name=connection.qualified_name
+ )
+ response = client.asset.save(to_create)
+ result = response.assets_created(asset_type=CustomEntity)[0]
+ yield result
+ delete_asset(client, guid=result.guid, asset_type=CustomEntity)
+
+
+def test_custom_entity(
+ client: AtlanClient, connection: Connection, custom_entity: CustomEntity
+):
+ assert custom_entity
+ assert custom_entity.guid
+ assert custom_entity.qualified_name
+ assert custom_entity.name == CUSTOM_ENTITY_NAME
+ assert custom_entity.connection_qualified_name == connection.qualified_name
+ assert custom_entity.connector_name == AtlanConnectorType.CUSTOM.value
+
+
+@pytest.mark.order(after="test_custom_entity")
+def test_delete_custom_entity(
+ client: AtlanClient,
+ connection: Connection,
+ custom_entity: CustomEntity,
+):
+ response = client.asset.delete_by_guid(custom_entity.guid)
+ assert response
+ assert not response.assets_created(asset_type=CustomEntity)
+ assert not response.assets_updated(asset_type=CustomEntity)
+ deleted = response.assets_deleted(asset_type=CustomEntity)
+ assert deleted
+ assert len(deleted) == 1
+ assert deleted[0].guid == custom_entity.guid
+ assert deleted[0].qualified_name == custom_entity.qualified_name
+ assert deleted[0].delete_handler == "SOFT"
+ assert deleted[0].status == EntityStatus.DELETED
+
+
+@pytest.mark.order(after="test_delete_custom_entity")
+def test_restore_custom_entity(
+ client: AtlanClient,
+ connection: Connection,
+ custom_entity: CustomEntity,
+):
+ assert custom_entity.qualified_name
+ assert client.asset.restore(
+ asset_type=CustomEntity, qualified_name=custom_entity.qualified_name
+ )
+ assert custom_entity.qualified_name
+ restored = client.asset.get_by_qualified_name(
+ asset_type=CustomEntity, qualified_name=custom_entity.qualified_name
+ )
+ assert restored
+ assert restored.guid == custom_entity.guid
+ assert restored.qualified_name == custom_entity.qualified_name
+ assert restored.status == EntityStatus.ACTIVE
diff --git a/tests/unit/model/constants.py b/tests/unit/model/constants.py
index 28f9a43b3..88f4fb821 100644
--- a/tests/unit/model/constants.py
+++ b/tests/unit/model/constants.py
@@ -237,4 +237,10 @@
SUPERSET_DATASET_QUALIFIED_NAME = (
f"{SUPERSET_DASHBOARD_QUALIFIED_NAME}/{SUPERSET_DATASET_NAME}"
)
+CUSTOM_CONNECTION_QUALIFIED_NAME = "default/custom/123456789"
+CUSTOM_CONNECTOR_TYPE = "custom"
+CUSTOM_ENTITY_NAME = "test-custom-entity"
+CUSTOM_ENTITY_QUALIFIED_NAME = (
+ f"{CUSTOM_CONNECTION_QUALIFIED_NAME}/{CUSTOM_ENTITY_NAME}"
+)
PROCEDURE_NAME = "test-procedure"
diff --git a/tests/unit/model/custom_entity_test.py b/tests/unit/model/custom_entity_test.py
new file mode 100644
index 000000000..bc49a322b
--- /dev/null
+++ b/tests/unit/model/custom_entity_test.py
@@ -0,0 +1,69 @@
+import pytest
+
+from pyatlan.model.assets import CustomEntity
+from tests.unit.model.constants import (
+ CUSTOM_ENTITY_NAME,
+ CUSTOM_ENTITY_QUALIFIED_NAME,
+ CUSTOM_CONNECTION_QUALIFIED_NAME,
+ CUSTOM_CONNECTOR_TYPE,
+)
+
+
+@pytest.mark.parametrize(
+ "name, connection_qualified_name, message",
+ [
+ (None, "connection/name", "name is required"),
+ (CUSTOM_ENTITY_NAME, None, "connection_qualified_name is required"),
+ ],
+)
+def test_create_with_missing_parameters_raise_value_error(
+ name: str, connection_qualified_name: str, message: str
+):
+ with pytest.raises(ValueError, match=message):
+ CustomEntity.create(
+ name=name, connection_qualified_name=connection_qualified_name
+ )
+
+
+def test_create():
+ sut = CustomEntity.creator(
+ name=CUSTOM_ENTITY_NAME,
+ connection_qualified_name=CUSTOM_CONNECTION_QUALIFIED_NAME,
+ )
+
+ assert sut.name == CUSTOM_ENTITY_NAME
+ assert sut.connection_qualified_name == CUSTOM_CONNECTION_QUALIFIED_NAME
+ assert sut.qualified_name == CUSTOM_ENTITY_QUALIFIED_NAME
+ assert sut.connector_name == CUSTOM_CONNECTOR_TYPE
+
+
+@pytest.mark.parametrize(
+ "qualified_name, name, message",
+ [
+ (None, CUSTOM_CONNECTION_QUALIFIED_NAME, "qualified_name is required"),
+ (CUSTOM_ENTITY_NAME, None, "name is required"),
+ ],
+)
+def test_create_for_modification_with_invalid_parameter_raises_value_error(
+ qualified_name: str, name: str, message: str
+):
+ with pytest.raises(ValueError, match=message):
+ CustomEntity.create_for_modification(qualified_name=qualified_name, name=name)
+
+
+def test_create_for_modification():
+ sut = CustomEntity.create_for_modification(
+ qualified_name=CUSTOM_ENTITY_QUALIFIED_NAME, name=CUSTOM_ENTITY_NAME
+ )
+
+ assert sut.qualified_name == CUSTOM_ENTITY_QUALIFIED_NAME
+ assert sut.name == CUSTOM_ENTITY_NAME
+
+
+def test_trim_to_required():
+ sut = CustomEntity.create_for_modification(
+ name=CUSTOM_ENTITY_NAME, qualified_name=CUSTOM_ENTITY_QUALIFIED_NAME
+ ).trim_to_required()
+
+ assert sut.name == CUSTOM_ENTITY_NAME
+ assert sut.qualified_name == CUSTOM_ENTITY_QUALIFIED_NAME
diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py
index c296d29bb..649fe839e 100644
--- a/tests/unit/test_model.py
+++ b/tests/unit/test_model.py
@@ -44,6 +44,7 @@
AuthPolicy,
Column,
ColumnProcess,
+ CustomEntity,
Database,
DbtMetric,
DbtModel,
@@ -383,6 +384,7 @@
"Optional[SupersetDashboard]": SupersetDashboard(),
"Optional[List[SupersetChart]]": [SupersetChart()],
"Optional[List[SupersetDataset]]": [SupersetDataset()],
+ "Optional[CustomEntity]": CustomEntity(),
"Optional[List[ModeCollection]]": [ModeCollection()],
"Optional[List[ModeQuery]]": [ModeQuery()],
"Optional[List[ModeChart]]": [ModeChart()],
From 865328cb5198fcc5f56c435fb341237dc57df12c Mon Sep 17 00:00:00 2001
From: rupali-atlan <rupali.kishore@atlan.com>
Date: Fri, 10 Jan 2025 19:12:06 +0530
Subject: [PATCH 2/6] did formatting via running black .
---
pyatlan/model/assets/custom_entity.py | 2 --
tests/integration/custom_asset_test.py | 4 +---
2 files changed, 1 insertion(+), 5 deletions(-)
diff --git a/pyatlan/model/assets/custom_entity.py b/pyatlan/model/assets/custom_entity.py
index 2630aa597..728d66cd2 100644
--- a/pyatlan/model/assets/custom_entity.py
+++ b/pyatlan/model/assets/custom_entity.py
@@ -140,7 +140,6 @@ class Attributes(Custom.Attributes):
default=None, description=""
) # relationship
-
@classmethod
@init_guid
def create(
@@ -158,7 +157,6 @@ def create(
),
)
-
attributes: CustomEntity.Attributes = Field(
default_factory=lambda: CustomEntity.Attributes(),
description=(
diff --git a/tests/integration/custom_asset_test.py b/tests/integration/custom_asset_test.py
index fa84bcc91..4cb7fc917 100644
--- a/tests/integration/custom_asset_test.py
+++ b/tests/integration/custom_asset_test.py
@@ -3,9 +3,7 @@
import pytest
from pyatlan.client.atlan import AtlanClient
-from pyatlan.model.assets import (
- Connection, CustomEntity
-)
+from pyatlan.model.assets import Connection, CustomEntity
from pyatlan.model.enums import (
AtlanConnectorType,
EntityStatus,
From 6ad021489ad29fb2c197522fc2da46c0963b987a Mon Sep 17 00:00:00 2001
From: rupali-atlan <rupali.kishore@atlan.com>
Date: Mon, 13 Jan 2025 12:14:06 +0530
Subject: [PATCH 3/6] used creator/updater methods
---
.../methods/asset/custom_entity.jinja2 | 2 +-
.../methods/attribute/custom_entity.jinja2 | 2 +-
pyatlan/model/assets/custom_entity.py | 4 +-
pyatlan/model/structs.py | 109 ++++++++----------
tests/unit/model/custom_entity_test.py | 16 +--
5 files changed, 63 insertions(+), 70 deletions(-)
diff --git a/pyatlan/generator/templates/methods/asset/custom_entity.jinja2 b/pyatlan/generator/templates/methods/asset/custom_entity.jinja2
index f78fb35d8..fb8f3e972 100644
--- a/pyatlan/generator/templates/methods/asset/custom_entity.jinja2
+++ b/pyatlan/generator/templates/methods/asset/custom_entity.jinja2
@@ -5,7 +5,7 @@
validate_required_fields(
["name", "connection_qualified_name"], [name, connection_qualified_name]
)
- attributes = CustomEntity.Attributes.create(
+ attributes = CustomEntity.Attributes.creator(
name=name, connection_qualified_name=connection_qualified_name
)
return cls(attributes=attributes)
diff --git a/pyatlan/generator/templates/methods/attribute/custom_entity.jinja2 b/pyatlan/generator/templates/methods/attribute/custom_entity.jinja2
index be2a575e5..9addd7dbe 100644
--- a/pyatlan/generator/templates/methods/attribute/custom_entity.jinja2
+++ b/pyatlan/generator/templates/methods/attribute/custom_entity.jinja2
@@ -1,7 +1,7 @@
@classmethod
@init_guid
- def create(
+ def creator(
cls, *, name: str, connection_qualified_name: str
) -> CustomEntity.Attributes:
validate_required_fields(
diff --git a/pyatlan/model/assets/custom_entity.py b/pyatlan/model/assets/custom_entity.py
index 728d66cd2..779895d8a 100644
--- a/pyatlan/model/assets/custom_entity.py
+++ b/pyatlan/model/assets/custom_entity.py
@@ -22,7 +22,7 @@ def creator(cls, *, name: str, connection_qualified_name: str) -> CustomEntity:
validate_required_fields(
["name", "connection_qualified_name"], [name, connection_qualified_name]
)
- attributes = CustomEntity.Attributes.create(
+ attributes = CustomEntity.Attributes.creator(
name=name, connection_qualified_name=connection_qualified_name
)
return cls(attributes=attributes)
@@ -142,7 +142,7 @@ class Attributes(Custom.Attributes):
@classmethod
@init_guid
- def create(
+ def creator(
cls, *, name: str, connection_qualified_name: str
) -> CustomEntity.Attributes:
validate_required_fields(
diff --git a/pyatlan/model/structs.py b/pyatlan/model/structs.py
index 42df88751..43e321cbb 100644
--- a/pyatlan/model/structs.py
+++ b/pyatlan/model/structs.py
@@ -78,13 +78,11 @@ class AwsCloudWatchMetric(AtlanObject):
aws_cloud_watch_metric_scope: str = Field(description="")
-class KafkaTopicConsumption(AtlanObject):
+class Histogram(AtlanObject):
"""Description"""
- topic_name: Optional[str] = Field(default=None, description="")
- topic_partition: Optional[str] = Field(default=None, description="")
- topic_lag: Optional[int] = Field(default=None, description="")
- topic_current_offset: Optional[int] = Field(default=None, description="")
+ boundaries: Set[float] = Field(description="")
+ frequencies: Set[float] = Field(description="")
class Action(AtlanObject):
@@ -96,11 +94,13 @@ class Action(AtlanObject):
task_action_display_text: Optional[str] = Field(default=None, description="")
-class Histogram(AtlanObject):
+class KafkaTopicConsumption(AtlanObject):
"""Description"""
- boundaries: Set[float] = Field(description="")
- frequencies: Set[float] = Field(description="")
+ topic_name: Optional[str] = Field(default=None, description="")
+ topic_partition: Optional[str] = Field(default=None, description="")
+ topic_lag: Optional[int] = Field(default=None, description="")
+ topic_current_offset: Optional[int] = Field(default=None, description="")
class ColumnValueFrequencyMap(AtlanObject):
@@ -110,11 +110,38 @@ class ColumnValueFrequencyMap(AtlanObject):
column_value_frequency: Optional[int] = Field(default=None, description="")
-class SourceTagAttachmentValue(AtlanObject):
+class BadgeCondition(AtlanObject):
"""Description"""
- tag_attachment_key: Optional[str] = Field(default=None, description="")
- tag_attachment_value: Optional[str] = Field(default=None, description="")
+ @classmethod
+ def create(
+ cls,
+ *,
+ badge_condition_operator: BadgeComparisonOperator,
+ badge_condition_value: str,
+ badge_condition_colorhex: Union[BadgeConditionColor, str],
+ ) -> "BadgeCondition":
+ validate_required_fields(
+ [
+ "badge_condition_operator",
+ "badge_condition_value",
+ "badge_condition_colorhex",
+ ],
+ [badge_condition_operator, badge_condition_value, badge_condition_colorhex],
+ )
+ return cls(
+ badge_condition_operator=badge_condition_operator.value,
+ badge_condition_value=badge_condition_value,
+ badge_condition_colorhex=(
+ badge_condition_colorhex.value
+ if isinstance(badge_condition_colorhex, BadgeConditionColor)
+ else badge_condition_colorhex
+ ),
+ )
+
+ badge_condition_operator: Optional[str] = Field(default=None, description="")
+ badge_condition_value: Optional[str] = Field(default=None, description="")
+ badge_condition_colorhex: Optional[str] = Field(default=None, description="")
class SourceTagAttachment(AtlanObject):
@@ -261,40 +288,6 @@ def of(
)
-class BadgeCondition(AtlanObject):
- """Description"""
-
- @classmethod
- def create(
- cls,
- *,
- badge_condition_operator: BadgeComparisonOperator,
- badge_condition_value: str,
- badge_condition_colorhex: Union[BadgeConditionColor, str],
- ) -> "BadgeCondition":
- validate_required_fields(
- [
- "badge_condition_operator",
- "badge_condition_value",
- "badge_condition_colorhex",
- ],
- [badge_condition_operator, badge_condition_value, badge_condition_colorhex],
- )
- return cls(
- badge_condition_operator=badge_condition_operator.value,
- badge_condition_value=badge_condition_value,
- badge_condition_colorhex=(
- badge_condition_colorhex.value
- if isinstance(badge_condition_colorhex, BadgeConditionColor)
- else badge_condition_colorhex
- ),
- )
-
- badge_condition_operator: Optional[str] = Field(default=None, description="")
- badge_condition_value: Optional[str] = Field(default=None, description="")
- badge_condition_colorhex: Optional[str] = Field(default=None, description="")
-
-
class AzureTag(AtlanObject):
"""Description"""
@@ -323,13 +316,6 @@ class AwsTag(AtlanObject):
aws_tag_value: str = Field(description="")
-class GoogleTag(AtlanObject):
- """Description"""
-
- google_tag_key: str = Field(description="")
- google_tag_value: str = Field(description="")
-
-
class DbtMetricFilter(AtlanObject):
"""Description"""
@@ -341,6 +327,13 @@ class DbtMetricFilter(AtlanObject):
dbt_metric_filter_value: Optional[str] = Field(default=None, description="")
+class GoogleTag(AtlanObject):
+ """Description"""
+
+ google_tag_key: str = Field(description="")
+ google_tag_value: str = Field(description="")
+
+
class AuthPolicyValiditySchedule(AtlanObject):
"""Description"""
@@ -402,20 +395,20 @@ class SourceTagAttribute(AtlanObject):
AwsCloudWatchMetric.update_forward_refs()
-KafkaTopicConsumption.update_forward_refs()
+Histogram.update_forward_refs()
Action.update_forward_refs()
-Histogram.update_forward_refs()
+KafkaTopicConsumption.update_forward_refs()
ColumnValueFrequencyMap.update_forward_refs()
+BadgeCondition.update_forward_refs()
+
SourceTagAttachmentValue.update_forward_refs()
SourceTagAttachment.update_forward_refs()
-BadgeCondition.update_forward_refs()
-
AzureTag.update_forward_refs()
StarredDetails.update_forward_refs()
@@ -424,10 +417,10 @@ class SourceTagAttribute(AtlanObject):
AwsTag.update_forward_refs()
-GoogleTag.update_forward_refs()
-
DbtMetricFilter.update_forward_refs()
+GoogleTag.update_forward_refs()
+
AuthPolicyValiditySchedule.update_forward_refs()
MCRuleComparison.update_forward_refs()
diff --git a/tests/unit/model/custom_entity_test.py b/tests/unit/model/custom_entity_test.py
index bc49a322b..d7cd34bdd 100644
--- a/tests/unit/model/custom_entity_test.py
+++ b/tests/unit/model/custom_entity_test.py
@@ -16,16 +16,16 @@
(CUSTOM_ENTITY_NAME, None, "connection_qualified_name is required"),
],
)
-def test_create_with_missing_parameters_raise_value_error(
+def test_creator_with_missing_parameters_raise_value_error(
name: str, connection_qualified_name: str, message: str
):
with pytest.raises(ValueError, match=message):
- CustomEntity.create(
+ CustomEntity.creator(
name=name, connection_qualified_name=connection_qualified_name
)
-def test_create():
+def test_creator():
sut = CustomEntity.creator(
name=CUSTOM_ENTITY_NAME,
connection_qualified_name=CUSTOM_CONNECTION_QUALIFIED_NAME,
@@ -44,15 +44,15 @@ def test_create():
(CUSTOM_ENTITY_NAME, None, "name is required"),
],
)
-def test_create_for_modification_with_invalid_parameter_raises_value_error(
+def test_updater_modification_with_invalid_parameter_raises_value_error(
qualified_name: str, name: str, message: str
):
with pytest.raises(ValueError, match=message):
- CustomEntity.create_for_modification(qualified_name=qualified_name, name=name)
+ CustomEntity.updater(qualified_name=qualified_name, name=name)
-def test_create_for_modification():
- sut = CustomEntity.create_for_modification(
+def test_updater():
+ sut = CustomEntity.updater(
qualified_name=CUSTOM_ENTITY_QUALIFIED_NAME, name=CUSTOM_ENTITY_NAME
)
@@ -61,7 +61,7 @@ def test_create_for_modification():
def test_trim_to_required():
- sut = CustomEntity.create_for_modification(
+ sut = CustomEntity.updater(
name=CUSTOM_ENTITY_NAME, qualified_name=CUSTOM_ENTITY_QUALIFIED_NAME
).trim_to_required()
From e2a015692b9e21c8646c2a908967d652962db8bd Mon Sep 17 00:00:00 2001
From: rupali-atlan <rupali.kishore@atlan.com>
Date: Mon, 13 Jan 2025 12:34:15 +0530
Subject: [PATCH 4/6] added update, retrieve integration tests for custom
assets
---
tests/integration/custom_asset_test.py | 66 +++++++++++++++++++++++++-
tests/unit/model/custom_entity_test.py | 2 +-
2 files changed, 66 insertions(+), 2 deletions(-)
diff --git a/tests/integration/custom_asset_test.py b/tests/integration/custom_asset_test.py
index 4cb7fc917..f02c35ec7 100644
--- a/tests/integration/custom_asset_test.py
+++ b/tests/integration/custom_asset_test.py
@@ -1,11 +1,14 @@
from typing import Generator
import pytest
+from pyatlan.model.core import Announcement
from pyatlan.client.atlan import AtlanClient
from pyatlan.model.assets import Connection, CustomEntity
from pyatlan.model.enums import (
+ AnnouncementType,
AtlanConnectorType,
+ CertificateStatus,
EntityStatus,
)
from pyatlan.model.response import AssetMutationResponse
@@ -18,6 +21,12 @@
CONNECTOR_TYPE = AtlanConnectorType.CUSTOM
CUSTOM_ENTITY_NAME = f"{MODULE_NAME}-custom-entity"
+CERTIFICATE_STATUS = CertificateStatus.VERIFIED
+ANNOUNCEMENT_TITLE = "Python SDK testing."
+ANNOUNCEMENT_TYPE = AnnouncementType.INFORMATION
+CERTIFICATE_MESSAGE = "Automated testing of the Python SDK."
+ANNOUNCEMENT_MESSAGE = "Automated testing of the Python SDK."
+
response = block(AtlanClient(), AssetMutationResponse())
@@ -54,7 +63,6 @@ def test_custom_entity(
assert custom_entity.connection_qualified_name == connection.qualified_name
assert custom_entity.connector_name == AtlanConnectorType.CUSTOM.value
-
@pytest.mark.order(after="test_custom_entity")
def test_delete_custom_entity(
client: AtlanClient,
@@ -92,3 +100,59 @@ def test_restore_custom_entity(
assert restored.guid == custom_entity.guid
assert restored.qualified_name == custom_entity.qualified_name
assert restored.status == EntityStatus.ACTIVE
+
+def _update_cert_and_annoucement(client, asset, asset_type):
+ assert asset.name
+ assert asset.qualified_name
+
+ updated = client.asset.update_certificate(
+ name=asset.name,
+ asset_type=asset_type,
+ qualified_name=asset.qualified_name,
+ message=CERTIFICATE_MESSAGE,
+ certificate_status=CERTIFICATE_STATUS,
+ )
+ assert updated
+ assert updated.certificate_status == CERTIFICATE_STATUS
+ assert updated.certificate_status_message == CERTIFICATE_MESSAGE
+
+ updated = client.asset.update_announcement(
+ name=asset.name,
+ asset_type=asset_type,
+ qualified_name=asset.qualified_name,
+ announcement=Announcement(
+ announcement_type=ANNOUNCEMENT_TYPE,
+ announcement_title=ANNOUNCEMENT_TITLE,
+ announcement_message=ANNOUNCEMENT_MESSAGE,
+ ),
+ )
+ assert updated
+ assert updated.announcement_type == ANNOUNCEMENT_TYPE
+ assert updated.announcement_title == ANNOUNCEMENT_TITLE
+ assert updated.announcement_message == ANNOUNCEMENT_MESSAGE
+
+def test_update_custom_assets(
+ client: AtlanClient,
+ custom_entity: CustomEntity,
+):
+ _update_cert_and_annoucement(client, custom_entity, CustomEntity)
+
+def _retrieve_custom_assets(client, asset, asset_type):
+ retrieved = client.asset.get_by_guid(
+ asset.guid, asset_type=asset_type, ignore_relationships=False
+ )
+ assert retrieved
+ assert not retrieved.is_incomplete
+ assert retrieved.guid == asset.guid
+ assert retrieved.qualified_name == asset.qualified_name
+ assert retrieved.name == asset.name
+ assert retrieved.connector_name == AtlanConnectorType.DATAVERSE
+ assert retrieved.certificate_status == CERTIFICATE_STATUS
+ assert retrieved.certificate_status_message == CERTIFICATE_MESSAGE
+
+@pytest.mark.order(after="test_update_custom_assets")
+def test_retrieve_custom_assets(
+ client: AtlanClient,
+ custom_entity: CustomEntity,
+):
+ _retrieve_custom_assets(client, custom_entity, CustomEntity)
diff --git a/tests/unit/model/custom_entity_test.py b/tests/unit/model/custom_entity_test.py
index d7cd34bdd..d09201bb3 100644
--- a/tests/unit/model/custom_entity_test.py
+++ b/tests/unit/model/custom_entity_test.py
@@ -44,7 +44,7 @@ def test_creator():
(CUSTOM_ENTITY_NAME, None, "name is required"),
],
)
-def test_updater_modification_with_invalid_parameter_raises_value_error(
+def test_updater_with_invalid_parameter_raises_value_error(
qualified_name: str, name: str, message: str
):
with pytest.raises(ValueError, match=message):
From 3287f4e1583798ce58b97a85b78ad2f87f7fb58c Mon Sep 17 00:00:00 2001
From: Aryamanz29 <aryamanz29@gmail.com>
Date: Mon, 13 Jan 2025 16:55:40 +0530
Subject: [PATCH 5/6] [gen] Generated latest typedefs
---
docs/asset/customentity.rst | 10 +
docs/asset/powerbidataflowentitycolumn.rst | 10 +
docs/assets.rst | 3 +
pyatlan/model/assets/__init__.py | 7 +-
pyatlan/model/assets/__init__.pyi | 6 +-
pyatlan/model/assets/core/__init__.py | 2 +
.../model/assets/core/power_b_i_dataflow.py | 57 ++++++
.../core/power_b_i_dataflow_entity_column.py | 177 ++++++++++++++++++
.../model/assets/core/power_b_i_datasource.py | 19 ++
pyatlan/model/assets/custom_entity.py | 6 +-
pyatlan/model/structs.py | 38 ++--
11 files changed, 308 insertions(+), 27 deletions(-)
create mode 100644 docs/asset/customentity.rst
create mode 100644 docs/asset/powerbidataflowentitycolumn.rst
create mode 100644 pyatlan/model/assets/core/power_b_i_dataflow_entity_column.py
diff --git a/docs/asset/customentity.rst b/docs/asset/customentity.rst
new file mode 100644
index 000000000..0da34b148
--- /dev/null
+++ b/docs/asset/customentity.rst
@@ -0,0 +1,10 @@
+.. _customentity:
+
+CustomEntity
+============
+
+.. module:: pyatlan.model.assets
+ :no-index:
+
+.. autoclass:: CustomEntity
+ :members:
diff --git a/docs/asset/powerbidataflowentitycolumn.rst b/docs/asset/powerbidataflowentitycolumn.rst
new file mode 100644
index 000000000..006c3380c
--- /dev/null
+++ b/docs/asset/powerbidataflowentitycolumn.rst
@@ -0,0 +1,10 @@
+.. _powerbidataflowentitycolumn:
+
+PowerBIDataflowEntityColumn
+===========================
+
+.. module:: pyatlan.model.assets
+ :no-index:
+
+.. autoclass:: PowerBIDataflowEntityColumn
+ :members:
diff --git a/docs/assets.rst b/docs/assets.rst
index 1e164b773..42014b6d2 100644
--- a/docs/assets.rst
+++ b/docs/assets.rst
@@ -95,6 +95,8 @@ You can interact with all of the following different kinds of assets:
asset/cubedimension
asset/cubefield
asset/cubehierarchy
+ asset/custom
+ asset/customentity
asset/datacontract
asset/datadomain
asset/datamesh
@@ -204,6 +206,7 @@ You can interact with all of the following different kinds of assets:
asset/powerbicolumn
asset/powerbidashboard
asset/powerbidataflow
+ asset/powerbidataflowentitycolumn
asset/powerbidataset
asset/powerbidatasource
asset/powerbimeasure
diff --git a/pyatlan/model/assets/__init__.py b/pyatlan/model/assets/__init__.py
index f87a62701..2828f838c 100644
--- a/pyatlan/model/assets/__init__.py
+++ b/pyatlan/model/assets/__init__.py
@@ -64,10 +64,10 @@
"Function",
"TablePartition",
"Column",
- "SnowflakeStream",
"DatabricksUnityCatalogTag",
- "Database",
+ "SnowflakeStream",
"CalculationView",
+ "Database",
"Procedure",
"SnowflakeTag",
"CosmosMongoDB",
@@ -104,6 +104,7 @@
"PowerBIDashboard",
"PowerBIDataflow",
"PowerBIPage",
+ "PowerBIDataflowEntityColumn",
"SnowflakeDynamicTable",
"MongoDBCollection",
"DynamoDBSecondaryIndex",
@@ -111,8 +112,6 @@
"CosmosMongoDBAccount",
"CosmosMongoDBCollection",
"CosmosMongoDBDatabase",
- "Custom",
- "CustomEntity",
],
"task": ["Task"],
"data_set": ["DataSet"],
diff --git a/pyatlan/model/assets/__init__.pyi b/pyatlan/model/assets/__init__.pyi
index 6653f1fdf..8cd5bb825 100644
--- a/pyatlan/model/assets/__init__.pyi
+++ b/pyatlan/model/assets/__init__.pyi
@@ -61,10 +61,10 @@ __all__ = [
"Function",
"TablePartition",
"Column",
- "SnowflakeStream",
"DatabricksUnityCatalogTag",
- "Database",
+ "SnowflakeStream",
"CalculationView",
+ "Database",
"Procedure",
"SnowflakeTag",
"CosmosMongoDB",
@@ -101,6 +101,7 @@ __all__ = [
"PowerBIDashboard",
"PowerBIDataflow",
"PowerBIPage",
+ "PowerBIDataflowEntityColumn",
"SnowflakeDynamicTable",
"MongoDBCollection",
"DynamoDBSecondaryIndex",
@@ -440,6 +441,7 @@ from .core.power_b_i import PowerBI
from .core.power_b_i_column import PowerBIColumn
from .core.power_b_i_dashboard import PowerBIDashboard
from .core.power_b_i_dataflow import PowerBIDataflow
+from .core.power_b_i_dataflow_entity_column import PowerBIDataflowEntityColumn
from .core.power_b_i_dataset import PowerBIDataset
from .core.power_b_i_datasource import PowerBIDatasource
from .core.power_b_i_measure import PowerBIMeasure
diff --git a/pyatlan/model/assets/core/__init__.py b/pyatlan/model/assets/core/__init__.py
index 933105b6d..6eae08e17 100644
--- a/pyatlan/model/assets/core/__init__.py
+++ b/pyatlan/model/assets/core/__init__.py
@@ -77,6 +77,7 @@
from .power_b_i_column import PowerBIColumn
from .power_b_i_dashboard import PowerBIDashboard
from .power_b_i_dataflow import PowerBIDataflow
+from .power_b_i_dataflow_entity_column import PowerBIDataflowEntityColumn
from .power_b_i_dataset import PowerBIDataset
from .power_b_i_datasource import PowerBIDatasource
from .power_b_i_measure import PowerBIMeasure
@@ -211,6 +212,7 @@
PowerBIDashboard.Attributes.update_forward_refs(**localns)
PowerBIDataflow.Attributes.update_forward_refs(**localns)
PowerBIPage.Attributes.update_forward_refs(**localns)
+PowerBIDataflowEntityColumn.Attributes.update_forward_refs(**localns)
SnowflakeDynamicTable.Attributes.update_forward_refs(**localns)
MongoDBCollection.Attributes.update_forward_refs(**localns)
DynamoDBSecondaryIndex.Attributes.update_forward_refs(**localns)
diff --git a/pyatlan/model/assets/core/power_b_i_dataflow.py b/pyatlan/model/assets/core/power_b_i_dataflow.py
index 8fce0a0f1..cfd91d473 100644
--- a/pyatlan/model/assets/core/power_b_i_dataflow.py
+++ b/pyatlan/model/assets/core/power_b_i_dataflow.py
@@ -72,10 +72,20 @@ def __setattr__(self, name, value):
"""
TBC
"""
+ POWER_BI_DATAFLOW_ENTITY_COLUMNS: ClassVar[RelationField] = RelationField(
+ "powerBIDataflowEntityColumns"
+ )
+ """
+ TBC
+ """
TABLES: ClassVar[RelationField] = RelationField("tables")
"""
TBC
"""
+ POWER_BI_DATASOURCES: ClassVar[RelationField] = RelationField("powerBIDatasources")
+ """
+ TBC
+ """
POWER_BI_DATAFLOW_CHILDREN: ClassVar[RelationField] = RelationField(
"powerBIDataflowChildren"
)
@@ -98,7 +108,9 @@ def __setattr__(self, name, value):
"workspace",
"power_b_i_processes",
"datasets",
+ "power_b_i_dataflow_entity_columns",
"tables",
+ "power_b_i_datasources",
"power_b_i_dataflow_children",
"power_b_i_dataflow_parents",
]
@@ -211,6 +223,27 @@ def datasets(self, datasets: Optional[List[PowerBIDataset]]):
self.attributes = self.Attributes()
self.attributes.datasets = datasets
+ @property
+ def power_b_i_dataflow_entity_columns(
+ self,
+ ) -> Optional[List[PowerBIDataflowEntityColumn]]:
+ return (
+ None
+ if self.attributes is None
+ else self.attributes.power_b_i_dataflow_entity_columns
+ )
+
+ @power_b_i_dataflow_entity_columns.setter
+ def power_b_i_dataflow_entity_columns(
+ self,
+ power_b_i_dataflow_entity_columns: Optional[List[PowerBIDataflowEntityColumn]],
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_dataflow_entity_columns = (
+ power_b_i_dataflow_entity_columns
+ )
+
@property
def tables(self) -> Optional[List[PowerBITable]]:
return None if self.attributes is None else self.attributes.tables
@@ -221,6 +254,20 @@ def tables(self, tables: Optional[List[PowerBITable]]):
self.attributes = self.Attributes()
self.attributes.tables = tables
+ @property
+ def power_b_i_datasources(self) -> Optional[List[PowerBIDatasource]]:
+ return (
+ None if self.attributes is None else self.attributes.power_b_i_datasources
+ )
+
+ @power_b_i_datasources.setter
+ def power_b_i_datasources(
+ self, power_b_i_datasources: Optional[List[PowerBIDatasource]]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_datasources = power_b_i_datasources
+
@property
def power_b_i_dataflow_children(self) -> Optional[List[PowerBIDataflow]]:
return (
@@ -274,9 +321,17 @@ class Attributes(PowerBI.Attributes):
datasets: Optional[List[PowerBIDataset]] = Field(
default=None, description=""
) # relationship
+ power_b_i_dataflow_entity_columns: Optional[
+ List[PowerBIDataflowEntityColumn]
+ ] = Field(
+ default=None, description=""
+ ) # relationship
tables: Optional[List[PowerBITable]] = Field(
default=None, description=""
) # relationship
+ power_b_i_datasources: Optional[List[PowerBIDatasource]] = Field(
+ default=None, description=""
+ ) # relationship
power_b_i_dataflow_children: Optional[List[PowerBIDataflow]] = Field(
default=None, description=""
) # relationship
@@ -294,7 +349,9 @@ class Attributes(PowerBI.Attributes):
)
+from .power_b_i_dataflow_entity_column import PowerBIDataflowEntityColumn # noqa
from .power_b_i_dataset import PowerBIDataset # noqa
+from .power_b_i_datasource import PowerBIDatasource # noqa
from .power_b_i_table import PowerBITable # noqa
from .power_b_i_workspace import PowerBIWorkspace # noqa
from .process import Process # noqa
diff --git a/pyatlan/model/assets/core/power_b_i_dataflow_entity_column.py b/pyatlan/model/assets/core/power_b_i_dataflow_entity_column.py
new file mode 100644
index 000000000..d734366c9
--- /dev/null
+++ b/pyatlan/model/assets/core/power_b_i_dataflow_entity_column.py
@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 Atlan Pte. Ltd.
+
+
+from __future__ import annotations
+
+from typing import ClassVar, List, Optional
+
+from pydantic.v1 import Field, validator
+
+from pyatlan.model.fields.atlan_fields import KeywordField, RelationField, TextField
+
+from .power_b_i import PowerBI
+
+
+class PowerBIDataflowEntityColumn(PowerBI):
+ """Description"""
+
+ type_name: str = Field(default="PowerBIDataflowEntityColumn", allow_mutation=False)
+
+ @validator("type_name")
+ def validate_type_name(cls, v):
+ if v != "PowerBIDataflowEntityColumn":
+ raise ValueError("must be PowerBIDataflowEntityColumn")
+ return v
+
+ def __setattr__(self, name, value):
+ if name in PowerBIDataflowEntityColumn._convenience_properties:
+ return object.__setattr__(self, name, value)
+ super().__setattr__(name, value)
+
+ POWER_BI_DATAFLOW_ENTITY_NAME: ClassVar[TextField] = TextField(
+ "powerBIDataflowEntityName", "powerBIDataflowEntityName"
+ )
+ """
+ Unique name of the dataflow entity in which this dataflow entity column exists.
+ """
+ POWER_BI_WORKSPACE_QUALIFIED_NAME: ClassVar[KeywordField] = KeywordField(
+ "powerBIWorkspaceQualifiedName", "powerBIWorkspaceQualifiedName"
+ )
+ """
+ Unique name of the workspace in which this dataflow entity column exists.
+ """
+ POWER_BI_DATAFLOW_QUALIFIED_NAME: ClassVar[KeywordField] = KeywordField(
+ "powerBIDataflowQualifiedName", "powerBIDataflowQualifiedName"
+ )
+ """
+ Unique name of the dataflow in which this dataflow entity column exists.
+ """
+ POWER_BI_DATAFLOW_ENTITY_COLUMN_DATA_TYPE: ClassVar[KeywordField] = KeywordField(
+ "powerBIDataflowEntityColumnDataType", "powerBIDataflowEntityColumnDataType"
+ )
+ """
+ Data type of this dataflow entity column.
+ """
+
+ POWER_BI_DATAFLOW: ClassVar[RelationField] = RelationField("powerBIDataflow")
+ """
+ TBC
+ """
+
+ _convenience_properties: ClassVar[List[str]] = [
+ "power_b_i_dataflow_entity_name",
+ "power_b_i_workspace_qualified_name",
+ "power_b_i_dataflow_qualified_name",
+ "power_b_i_dataflow_entity_column_data_type",
+ "power_b_i_dataflow",
+ ]
+
+ @property
+ def power_b_i_dataflow_entity_name(self) -> Optional[str]:
+ return (
+ None
+ if self.attributes is None
+ else self.attributes.power_b_i_dataflow_entity_name
+ )
+
+ @power_b_i_dataflow_entity_name.setter
+ def power_b_i_dataflow_entity_name(
+ self, power_b_i_dataflow_entity_name: Optional[str]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_dataflow_entity_name = power_b_i_dataflow_entity_name
+
+ @property
+ def power_b_i_workspace_qualified_name(self) -> Optional[str]:
+ return (
+ None
+ if self.attributes is None
+ else self.attributes.power_b_i_workspace_qualified_name
+ )
+
+ @power_b_i_workspace_qualified_name.setter
+ def power_b_i_workspace_qualified_name(
+ self, power_b_i_workspace_qualified_name: Optional[str]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_workspace_qualified_name = (
+ power_b_i_workspace_qualified_name
+ )
+
+ @property
+ def power_b_i_dataflow_qualified_name(self) -> Optional[str]:
+ return (
+ None
+ if self.attributes is None
+ else self.attributes.power_b_i_dataflow_qualified_name
+ )
+
+ @power_b_i_dataflow_qualified_name.setter
+ def power_b_i_dataflow_qualified_name(
+ self, power_b_i_dataflow_qualified_name: Optional[str]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_dataflow_qualified_name = (
+ power_b_i_dataflow_qualified_name
+ )
+
+ @property
+ def power_b_i_dataflow_entity_column_data_type(self) -> Optional[str]:
+ return (
+ None
+ if self.attributes is None
+ else self.attributes.power_b_i_dataflow_entity_column_data_type
+ )
+
+ @power_b_i_dataflow_entity_column_data_type.setter
+ def power_b_i_dataflow_entity_column_data_type(
+ self, power_b_i_dataflow_entity_column_data_type: Optional[str]
+ ):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_dataflow_entity_column_data_type = (
+ power_b_i_dataflow_entity_column_data_type
+ )
+
+ @property
+ def power_b_i_dataflow(self) -> Optional[PowerBIDataflow]:
+ return None if self.attributes is None else self.attributes.power_b_i_dataflow
+
+ @power_b_i_dataflow.setter
+ def power_b_i_dataflow(self, power_b_i_dataflow: Optional[PowerBIDataflow]):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_dataflow = power_b_i_dataflow
+
+ class Attributes(PowerBI.Attributes):
+ power_b_i_dataflow_entity_name: Optional[str] = Field(
+ default=None, description=""
+ )
+ power_b_i_workspace_qualified_name: Optional[str] = Field(
+ default=None, description=""
+ )
+ power_b_i_dataflow_qualified_name: Optional[str] = Field(
+ default=None, description=""
+ )
+ power_b_i_dataflow_entity_column_data_type: Optional[str] = Field(
+ default=None, description=""
+ )
+ power_b_i_dataflow: Optional[PowerBIDataflow] = Field(
+ default=None, description=""
+ ) # relationship
+
+ attributes: PowerBIDataflowEntityColumn.Attributes = Field(
+ default_factory=lambda: PowerBIDataflowEntityColumn.Attributes(),
+ description=(
+ "Map of attributes in the instance and their values. "
+ "The specific keys of this map will vary by type, "
+ "so are described in the sub-types of this schema."
+ ),
+ )
+
+
+from .power_b_i_dataflow import PowerBIDataflow # noqa
diff --git a/pyatlan/model/assets/core/power_b_i_datasource.py b/pyatlan/model/assets/core/power_b_i_datasource.py
index ae2a96069..8accfcca2 100644
--- a/pyatlan/model/assets/core/power_b_i_datasource.py
+++ b/pyatlan/model/assets/core/power_b_i_datasource.py
@@ -36,6 +36,10 @@ def __setattr__(self, name, value):
Connection details of the datasource.
"""
+ POWER_BI_DATAFLOWS: ClassVar[RelationField] = RelationField("powerBIDataflows")
+ """
+ TBC
+ """
DATASETS: ClassVar[RelationField] = RelationField("datasets")
"""
TBC
@@ -43,6 +47,7 @@ def __setattr__(self, name, value):
_convenience_properties: ClassVar[List[str]] = [
"connection_details",
+ "power_b_i_dataflows",
"datasets",
]
@@ -56,6 +61,16 @@ def connection_details(self, connection_details: Optional[Dict[str, str]]):
self.attributes = self.Attributes()
self.attributes.connection_details = connection_details
+ @property
+ def power_b_i_dataflows(self) -> Optional[List[PowerBIDataflow]]:
+ return None if self.attributes is None else self.attributes.power_b_i_dataflows
+
+ @power_b_i_dataflows.setter
+ def power_b_i_dataflows(self, power_b_i_dataflows: Optional[List[PowerBIDataflow]]):
+ if self.attributes is None:
+ self.attributes = self.Attributes()
+ self.attributes.power_b_i_dataflows = power_b_i_dataflows
+
@property
def datasets(self) -> Optional[List[PowerBIDataset]]:
return None if self.attributes is None else self.attributes.datasets
@@ -70,6 +85,9 @@ class Attributes(PowerBI.Attributes):
connection_details: Optional[Dict[str, str]] = Field(
default=None, description=""
)
+ power_b_i_dataflows: Optional[List[PowerBIDataflow]] = Field(
+ default=None, description=""
+ ) # relationship
datasets: Optional[List[PowerBIDataset]] = Field(
default=None, description=""
) # relationship
@@ -84,4 +102,5 @@ class Attributes(PowerBI.Attributes):
)
+from .power_b_i_dataflow import PowerBIDataflow # noqa
from .power_b_i_dataset import PowerBIDataset # noqa
diff --git a/pyatlan/model/assets/custom_entity.py b/pyatlan/model/assets/custom_entity.py
index 779895d8a..98ed2e7fd 100644
--- a/pyatlan/model/assets/custom_entity.py
+++ b/pyatlan/model/assets/custom_entity.py
@@ -5,10 +5,12 @@
from __future__ import annotations
from typing import ClassVar, List, Optional
-from pyatlan.model.enums import AtlanConnectorType
+
from pydantic.v1 import Field, validator
-from pyatlan.utils import init_guid, validate_required_fields
+
+from pyatlan.model.enums import AtlanConnectorType
from pyatlan.model.fields.atlan_fields import RelationField
+from pyatlan.utils import init_guid, validate_required_fields
from .custom import Custom
diff --git a/pyatlan/model/structs.py b/pyatlan/model/structs.py
index d2c11a48b..633dd2f50 100644
--- a/pyatlan/model/structs.py
+++ b/pyatlan/model/structs.py
@@ -78,11 +78,13 @@ class AwsCloudWatchMetric(AtlanObject):
aws_cloud_watch_metric_scope: str = Field(description="")
-class Histogram(AtlanObject):
+class KafkaTopicConsumption(AtlanObject):
"""Description"""
- boundaries: Set[float] = Field(description="")
- frequencies: Set[float] = Field(description="")
+ topic_name: Optional[str] = Field(default=None, description="")
+ topic_partition: Optional[str] = Field(default=None, description="")
+ topic_lag: Optional[int] = Field(default=None, description="")
+ topic_current_offset: Optional[int] = Field(default=None, description="")
class Action(AtlanObject):
@@ -94,13 +96,11 @@ class Action(AtlanObject):
task_action_display_text: Optional[str] = Field(default=None, description="")
-class KafkaTopicConsumption(AtlanObject):
+class Histogram(AtlanObject):
"""Description"""
- topic_name: Optional[str] = Field(default=None, description="")
- topic_partition: Optional[str] = Field(default=None, description="")
- topic_lag: Optional[int] = Field(default=None, description="")
- topic_current_offset: Optional[int] = Field(default=None, description="")
+ boundaries: Set[float] = Field(description="")
+ frequencies: Set[float] = Field(description="")
class ColumnValueFrequencyMap(AtlanObject):
@@ -323,6 +323,13 @@ class AwsTag(AtlanObject):
aws_tag_value: str = Field(description="")
+class GoogleTag(AtlanObject):
+ """Description"""
+
+ google_tag_key: str = Field(description="")
+ google_tag_value: str = Field(description="")
+
+
class DbtMetricFilter(AtlanObject):
"""Description"""
@@ -334,13 +341,6 @@ class DbtMetricFilter(AtlanObject):
dbt_metric_filter_value: Optional[str] = Field(default=None, description="")
-class GoogleTag(AtlanObject):
- """Description"""
-
- google_tag_key: str = Field(description="")
- google_tag_value: str = Field(description="")
-
-
class AuthPolicyValiditySchedule(AtlanObject):
"""Description"""
@@ -402,11 +402,11 @@ class SourceTagAttribute(AtlanObject):
AwsCloudWatchMetric.update_forward_refs()
-Histogram.update_forward_refs()
+KafkaTopicConsumption.update_forward_refs()
Action.update_forward_refs()
-KafkaTopicConsumption.update_forward_refs()
+Histogram.update_forward_refs()
ColumnValueFrequencyMap.update_forward_refs()
@@ -424,10 +424,10 @@ class SourceTagAttribute(AtlanObject):
AwsTag.update_forward_refs()
-DbtMetricFilter.update_forward_refs()
-
GoogleTag.update_forward_refs()
+DbtMetricFilter.update_forward_refs()
+
AuthPolicyValiditySchedule.update_forward_refs()
MCRuleComparison.update_forward_refs()
From b5a90a03e2ac44077d3e85275e2962624ec896ef Mon Sep 17 00:00:00 2001
From: Aryamanz29 <aryamanz29@gmail.com>
Date: Mon, 13 Jan 2025 17:07:07 +0530
Subject: [PATCH 6/6] [test/fix] Used `custom` connector
---
tests/integration/custom_asset_test.py | 9 +++++++--
1 file changed, 7 insertions(+), 2 deletions(-)
diff --git a/tests/integration/custom_asset_test.py b/tests/integration/custom_asset_test.py
index f02c35ec7..380f90cd8 100644
--- a/tests/integration/custom_asset_test.py
+++ b/tests/integration/custom_asset_test.py
@@ -1,10 +1,10 @@
from typing import Generator
import pytest
-from pyatlan.model.core import Announcement
from pyatlan.client.atlan import AtlanClient
from pyatlan.model.assets import Connection, CustomEntity
+from pyatlan.model.core import Announcement
from pyatlan.model.enums import (
AnnouncementType,
AtlanConnectorType,
@@ -63,6 +63,7 @@ def test_custom_entity(
assert custom_entity.connection_qualified_name == connection.qualified_name
assert custom_entity.connector_name == AtlanConnectorType.CUSTOM.value
+
@pytest.mark.order(after="test_custom_entity")
def test_delete_custom_entity(
client: AtlanClient,
@@ -101,6 +102,7 @@ def test_restore_custom_entity(
assert restored.qualified_name == custom_entity.qualified_name
assert restored.status == EntityStatus.ACTIVE
+
def _update_cert_and_annoucement(client, asset, asset_type):
assert asset.name
assert asset.qualified_name
@@ -131,12 +133,14 @@ def _update_cert_and_annoucement(client, asset, asset_type):
assert updated.announcement_title == ANNOUNCEMENT_TITLE
assert updated.announcement_message == ANNOUNCEMENT_MESSAGE
+
def test_update_custom_assets(
client: AtlanClient,
custom_entity: CustomEntity,
):
_update_cert_and_annoucement(client, custom_entity, CustomEntity)
+
def _retrieve_custom_assets(client, asset, asset_type):
retrieved = client.asset.get_by_guid(
asset.guid, asset_type=asset_type, ignore_relationships=False
@@ -146,10 +150,11 @@ def _retrieve_custom_assets(client, asset, asset_type):
assert retrieved.guid == asset.guid
assert retrieved.qualified_name == asset.qualified_name
assert retrieved.name == asset.name
- assert retrieved.connector_name == AtlanConnectorType.DATAVERSE
+ assert retrieved.connector_name == AtlanConnectorType.CUSTOM
assert retrieved.certificate_status == CERTIFICATE_STATUS
assert retrieved.certificate_status_message == CERTIFICATE_MESSAGE
+
@pytest.mark.order(after="test_update_custom_assets")
def test_retrieve_custom_assets(
client: AtlanClient,