From 357685b77d8d2c925bf6469b218932a3b9714ba1 Mon Sep 17 00:00:00 2001 From: Ernest Hill Date: Tue, 7 Jan 2025 14:09:45 +0300 Subject: [PATCH 1/6] Initial version --- pyatlan/model/assets/core/procedure.py | 102 ++++++++++++++++++++++++- tests/unit/model/constants.py | 1 + tests/unit/model/procedure_test.py | 82 ++++++++++++++++++++ 3 files changed, 184 insertions(+), 1 deletion(-) create mode 100644 tests/unit/model/procedure_test.py diff --git a/pyatlan/model/assets/core/procedure.py b/pyatlan/model/assets/core/procedure.py index 25f93eb74..95c571edc 100644 --- a/pyatlan/model/assets/core/procedure.py +++ b/pyatlan/model/assets/core/procedure.py @@ -4,18 +4,68 @@ from __future__ import annotations -from typing import ClassVar, List, Optional +from typing import ClassVar, List, Optional, overload from pydantic.v1 import Field, validator from pyatlan.model.fields.atlan_fields import RelationField, TextField +from pyatlan.utils import init_guid, validate_required_fields +from ...enums import AtlanConnectorType from .s_q_l import SQL class Procedure(SQL): """Description""" + @overload + @classmethod + def creator( + cls, + *, + name: str, + schema_qualified_name: str, + definition: str, + ) -> Procedure: ... + + @overload + @classmethod + def creator( + cls, + *, + name: str, + schema_qualified_name: str, + schema_name: str, + database_name: str, + database_qualified_name: str, + connection_qualified_name: str, + definition: str, + ) -> Procedure: ... + + @classmethod + @init_guid + def creator( + cls, + *, + name: str, + definition: str, + schema_qualified_name: str, + schema_name: Optional[str] = None, + database_name: Optional[str] = None, + database_qualified_name: Optional[str] = None, + connection_qualified_name: Optional[str] = None, + ) -> Procedure: + attributes = Procedure.Attributes.create( + name=name, + definition=definition, + schema_qualified_name=schema_qualified_name, + schema_name=schema_name, + database_name=database_name, + database_qualified_name=database_qualified_name, + connection_qualified_name=connection_qualified_name, + ) + return cls(attributes=attributes) + type_name: str = Field(default="Procedure", allow_mutation=False) @validator("type_name") @@ -70,6 +120,56 @@ class Attributes(SQL.Attributes): default=None, description="" ) # relationship + @classmethod + @init_guid + def create( + cls, + *, + name: str, + definition: str, + schema_qualified_name: Optional[str] = None, + schema_name: Optional[str] = None, + database_name: Optional[str] = None, + database_qualified_name: Optional[str] = None, + connection_qualified_name: Optional[str] = None, + ) -> Procedure.Attributes: + validate_required_fields( + ["name", "definition", "schema_qualified_name"], + [name, definition, schema_qualified_name], + ) + assert schema_qualified_name # noqa: S101 + if connection_qualified_name: + connector_name = AtlanConnectorType.get_connector_name( + connection_qualified_name + ) + else: + connection_qn, connector_name = AtlanConnectorType.get_connector_name( + schema_qualified_name, "schema_qualified_name", 5 + ) + + fields = schema_qualified_name.split("/") + qualified_name = f"{schema_qualified_name}/_procedures_/{name}" + connection_qualified_name = connection_qualified_name or connection_qn + database_name = database_name or fields[3] + schema_name = schema_name or fields[4] + database_qualified_name = ( + database_qualified_name + or f"{connection_qualified_name}/{database_name}" + ) + + return Procedure.Attributes( + name=name, + definition=definition, + qualified_name=qualified_name, + database_name=database_name, + database_qualified_name=database_qualified_name, + schema_name=schema_name, + schema_qualified_name=schema_qualified_name, + atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name), + connector_name=connector_name, + connection_qualified_name=connection_qualified_name, + ) + attributes: Procedure.Attributes = Field( default_factory=lambda: Procedure.Attributes(), description=( diff --git a/tests/unit/model/constants.py b/tests/unit/model/constants.py index 93fed9ae2..28f9a43b3 100644 --- a/tests/unit/model/constants.py +++ b/tests/unit/model/constants.py @@ -237,3 +237,4 @@ SUPERSET_DATASET_QUALIFIED_NAME = ( f"{SUPERSET_DASHBOARD_QUALIFIED_NAME}/{SUPERSET_DATASET_NAME}" ) +PROCEDURE_NAME = "test-procedure" diff --git a/tests/unit/model/procedure_test.py b/tests/unit/model/procedure_test.py new file mode 100644 index 000000000..42c9cf280 --- /dev/null +++ b/tests/unit/model/procedure_test.py @@ -0,0 +1,82 @@ +import pytest + +from pyatlan.model.assets import Procedure +from tests.unit.model.constants import ( + CONNECTION_QUALIFIED_NAME, + CONNECTOR_TYPE, + DATABASE_NAME, + DATABASE_QUALIFIED_NAME, + PROCEDURE_NAME, + SCHEMA_NAME, + SCHEMA_QUALIFIED_NAME, +) + +DEFINITION = """ +BEGIN +insert into `atlanhq.testing_lineage.INSTACART_ALCOHOL_ORDER_TIME_copy` +select * from `atlanhq.testing_lineage.INSTACART_ALCOHOL_ORDER_TIME`; +END +""" + + +@pytest.mark.parametrize( + "name, definition, schema_qualified_name, message", + [ + (None, DEFINITION, SCHEMA_QUALIFIED_NAME, "name is required"), + (PROCEDURE_NAME, None, SCHEMA_QUALIFIED_NAME, "definition is required"), + (PROCEDURE_NAME, DEFINITION, None, "schema_qualified_name is required"), + ], +) +def test_create_with_missing_parameters_raise_value_error( + name: str, definition: str, schema_qualified_name: str, message: str +): + with pytest.raises(ValueError, match=message): + Procedure.create( + name=name, + definition=definition, + schema_qualified_name=schema_qualified_name, + ) + + +def test_creator(): + sut = Procedure.create( + name=PROCEDURE_NAME, + definition=DEFINITION, + schema_qualified_name=SCHEMA_QUALIFIED_NAME, + ) + + assert sut.name == PROCEDURE_NAME + assert sut.database_name == DATABASE_NAME + assert sut.connection_qualified_name == CONNECTION_QUALIFIED_NAME + assert sut.database_qualified_name == DATABASE_QUALIFIED_NAME + assert ( + sut.qualified_name == f"{SCHEMA_QUALIFIED_NAME}/_procedures_/{PROCEDURE_NAME}" + ) + assert sut.schema_qualified_name == SCHEMA_QUALIFIED_NAME + assert sut.schema_name == SCHEMA_NAME + assert sut.connector_name == CONNECTOR_TYPE + assert sut.atlan_schema.qualified_name == SCHEMA_QUALIFIED_NAME + + +def test_overload_creator(): + sut = Procedure.creator( + name=PROCEDURE_NAME, + definition=DEFINITION, + schema_qualified_name=SCHEMA_QUALIFIED_NAME, + schema_name=SCHEMA_NAME, + database_name=DATABASE_NAME, + database_qualified_name=DATABASE_QUALIFIED_NAME, + connection_qualified_name=CONNECTION_QUALIFIED_NAME, + ) + + assert sut.name == PROCEDURE_NAME + assert sut.database_name == DATABASE_NAME + assert sut.connection_qualified_name == CONNECTION_QUALIFIED_NAME + assert sut.database_qualified_name == DATABASE_QUALIFIED_NAME + assert ( + sut.qualified_name == f"{SCHEMA_QUALIFIED_NAME}/_procedures_/{PROCEDURE_NAME}" + ) + assert sut.schema_qualified_name == SCHEMA_QUALIFIED_NAME + assert sut.schema_name == SCHEMA_NAME + assert sut.connector_name == CONNECTOR_TYPE + assert sut.atlan_schema.qualified_name == SCHEMA_QUALIFIED_NAME From 5bd34502ebb972ab52cba7a205370a811068f7f7 Mon Sep 17 00:00:00 2001 From: Ernest Hill Date: Tue, 7 Jan 2025 14:12:40 +0300 Subject: [PATCH 2/6] Correct typo. --- tests/unit/model/procedure_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/model/procedure_test.py b/tests/unit/model/procedure_test.py index 42c9cf280..088b7da58 100644 --- a/tests/unit/model/procedure_test.py +++ b/tests/unit/model/procedure_test.py @@ -39,7 +39,7 @@ def test_create_with_missing_parameters_raise_value_error( def test_creator(): - sut = Procedure.create( + sut = Procedure.creator( name=PROCEDURE_NAME, definition=DEFINITION, schema_qualified_name=SCHEMA_QUALIFIED_NAME, From 8e187b65512124d522c8e95a322d208e5af8dd35 Mon Sep 17 00:00:00 2001 From: Ernest Hill Date: Tue, 7 Jan 2025 15:11:09 +0300 Subject: [PATCH 3/6] After model regeneration --- docs/assets.rst | 3 + .../templates/methods/asset/procedure.jinja2 | 48 +++++++ .../templates/methods/asset/s3_object.jinja2 | 24 +++- .../methods/attribute/procedure.jinja2 | 50 ++++++++ .../methods/attribute/s3_object.jinja2 | 29 ++++- pyatlan/model/assets/__init__.py | 5 +- pyatlan/model/assets/__init__.pyi | 8 +- pyatlan/model/assets/core/__init__.py | 2 +- pyatlan/model/assets/core/procedure.py | 2 +- pyatlan/model/structs.py | 120 +++++++++--------- 10 files changed, 222 insertions(+), 69 deletions(-) create mode 100644 pyatlan/generator/templates/methods/asset/procedure.jinja2 create mode 100644 pyatlan/generator/templates/methods/attribute/procedure.jinja2 diff --git a/docs/assets.rst b/docs/assets.rst index 7e87e694b..1e164b773 100644 --- a/docs/assets.rst +++ b/docs/assets.rst @@ -105,6 +105,9 @@ You can interact with all of the following different kinds of assets: asset/datastudioasset asset/database asset/databricksunitycatalogtag + asset/dataverse + asset/dataverseattribute + asset/dataverseentity asset/dbt asset/dbtcolumnprocess asset/dbtmetric diff --git a/pyatlan/generator/templates/methods/asset/procedure.jinja2 b/pyatlan/generator/templates/methods/asset/procedure.jinja2 new file mode 100644 index 000000000..f4dcae109 --- /dev/null +++ b/pyatlan/generator/templates/methods/asset/procedure.jinja2 @@ -0,0 +1,48 @@ + + @overload + @classmethod + def creator( + cls, + *, + name: str, + schema_qualified_name: str, + definition: str, + ) -> Procedure: ... + + @overload + @classmethod + def creator( + cls, + *, + name: str, + schema_qualified_name: str, + schema_name: str, + database_name: str, + database_qualified_name: str, + connection_qualified_name: str, + definition: str, + ) -> Procedure: ... + + @classmethod + @init_guid + def creator( + cls, + *, + name: str, + definition: str, + schema_qualified_name: str, + schema_name: Optional[str] = None, + database_name: Optional[str] = None, + database_qualified_name: Optional[str] = None, + connection_qualified_name: Optional[str] = None, + ) -> Procedure: + attributes = Procedure.Attributes.create( + name=name, + definition=definition, + schema_qualified_name=schema_qualified_name, + schema_name=schema_name, + database_name=database_name, + database_qualified_name=database_qualified_name, + connection_qualified_name=connection_qualified_name, + ) + return cls(attributes=attributes) diff --git a/pyatlan/generator/templates/methods/asset/s3_object.jinja2 b/pyatlan/generator/templates/methods/asset/s3_object.jinja2 index a6f56cacb..abd40cc85 100644 --- a/pyatlan/generator/templates/methods/asset/s3_object.jinja2 +++ b/pyatlan/generator/templates/methods/asset/s3_object.jinja2 @@ -7,6 +7,7 @@ name: str, connection_qualified_name: str, aws_arn: str, + s3_bucket_name: str, s3_bucket_qualified_name: str, ) -> S3Object: validate_required_fields( @@ -14,14 +15,22 @@ "name", "connection_qualified_name", "aws_arn", + "s3_bucket_name", "s3_bucket_qualified_name", ], - [name, connection_qualified_name, aws_arn, s3_bucket_qualified_name], + [ + name, + connection_qualified_name, + aws_arn, + s3_bucket_name, + s3_bucket_qualified_name, + ], ) attributes = S3Object.Attributes.create( name=name, connection_qualified_name=connection_qualified_name, aws_arn=aws_arn, + s3_bucket_name=s3_bucket_name, s3_bucket_qualified_name=s3_bucket_qualified_name, ) return cls(attributes=attributes) @@ -34,6 +43,7 @@ name: str, connection_qualified_name: str, aws_arn: str, + s3_bucket_name: str, s3_bucket_qualified_name: str, ) -> S3Object: warn( @@ -48,6 +58,7 @@ name=name, connection_qualified_name=connection_qualified_name, aws_arn=aws_arn, + s3_bucket_name=s3_bucket_name, s3_bucket_qualified_name=s3_bucket_qualified_name, ) @@ -59,6 +70,7 @@ name: str, connection_qualified_name: str, prefix: str, + s3_bucket_name: str, s3_bucket_qualified_name: str, ) -> S3Object: validate_required_fields( @@ -66,14 +78,22 @@ "name", "connection_qualified_name", "prefix", + "s3_bucket_name", "s3_bucket_qualified_name", ], - [name, connection_qualified_name, prefix, s3_bucket_qualified_name], + [ + name, + connection_qualified_name, + prefix, + s3_bucket_name, + s3_bucket_qualified_name, + ], ) attributes = S3Object.Attributes.create_with_prefix( name=name, connection_qualified_name=connection_qualified_name, prefix=prefix, + s3_bucket_name=s3_bucket_name, s3_bucket_qualified_name=s3_bucket_qualified_name, ) return cls(attributes=attributes) diff --git a/pyatlan/generator/templates/methods/attribute/procedure.jinja2 b/pyatlan/generator/templates/methods/attribute/procedure.jinja2 new file mode 100644 index 000000000..5a1fe1d93 --- /dev/null +++ b/pyatlan/generator/templates/methods/attribute/procedure.jinja2 @@ -0,0 +1,50 @@ + + @classmethod + @init_guid + def create( + cls, + *, + name: str, + definition: str, + schema_qualified_name: Optional[str] = None, + schema_name: Optional[str] = None, + database_name: Optional[str] = None, + database_qualified_name: Optional[str] = None, + connection_qualified_name: Optional[str] = None, + ) -> Procedure.Attributes: + validate_required_fields( + ["name", "definition", "schema_qualified_name"], + [name, definition, schema_qualified_name], + ) + assert schema_qualified_name # noqa: S101 + if connection_qualified_name: + connector_name = AtlanConnectorType.get_connector_name( + connection_qualified_name + ) + else: + connection_qn, connector_name = AtlanConnectorType.get_connector_name( + schema_qualified_name, "schema_qualified_name", 5 + ) + + fields = schema_qualified_name.split("/") + qualified_name = f"{schema_qualified_name}/_procedures_/{name}" + connection_qualified_name = connection_qualified_name or connection_qn + database_name = database_name or fields[3] + schema_name = schema_name or fields[4] + database_qualified_name = ( + database_qualified_name + or f"{connection_qualified_name}/{database_name}" + ) + + return Procedure.Attributes( + name=name, + definition=definition, + qualified_name=qualified_name, + database_name=database_name, + database_qualified_name=database_qualified_name, + schema_name=schema_name, + schema_qualified_name=schema_qualified_name, + atlan_schema=Schema.ref_by_qualified_name(schema_qualified_name), + connector_name=connector_name, + connection_qualified_name=connection_qualified_name, + ) diff --git a/pyatlan/generator/templates/methods/attribute/s3_object.jinja2 b/pyatlan/generator/templates/methods/attribute/s3_object.jinja2 index 8c0fe564d..1ff67a8a9 100644 --- a/pyatlan/generator/templates/methods/attribute/s3_object.jinja2 +++ b/pyatlan/generator/templates/methods/attribute/s3_object.jinja2 @@ -7,11 +7,24 @@ name: str, connection_qualified_name: str, aws_arn: str, + s3_bucket_name: str, s3_bucket_qualified_name: str, ) -> S3Object.Attributes: validate_required_fields( - ["name", "connection_qualified_name", "aws_arn", "s3_bucket_qualified_name"], - [name, connection_qualified_name, aws_arn, s3_bucket_qualified_name], + [ + "name", + "connection_qualified_name", + "aws_arn", + "s3_bucket_name", + "s3_bucket_qualified_name", + ], + [ + name, + connection_qualified_name, + aws_arn, + s3_bucket_name, + s3_bucket_qualified_name, + ], ) fields = connection_qualified_name.split("/") if len(fields) != 3: @@ -30,6 +43,7 @@ connection_qualified_name=connection_qualified_name, qualified_name=f"{connection_qualified_name}/{aws_arn}", connector_name=connector_type.value, + s3_bucket_name=s3_bucket_name, s3_bucket_qualified_name=s3_bucket_qualified_name, bucket=S3Bucket.ref_by_qualified_name(s3_bucket_qualified_name), ) @@ -42,6 +56,7 @@ name: str, connection_qualified_name: str, prefix: str, + s3_bucket_name: str, s3_bucket_qualified_name: str, ) -> S3Object.Attributes: validate_required_fields( @@ -49,9 +64,16 @@ "name", "connection_qualified_name", "prefix", + "s3_bucket_name", "s3_bucket_qualified_name", ], - [name, connection_qualified_name, prefix, s3_bucket_qualified_name], + [ + name, + connection_qualified_name, + prefix, + s3_bucket_name, + s3_bucket_qualified_name, + ], ) fields = connection_qualified_name.split("/") if len(fields) != 3: @@ -71,6 +93,7 @@ connection_qualified_name=connection_qualified_name, qualified_name=f"{connection_qualified_name}/{object_key}", connector_name=connector_type.value, + s3_bucket_name=s3_bucket_name, s3_bucket_qualified_name=s3_bucket_qualified_name, bucket=S3Bucket.ref_by_qualified_name(s3_bucket_qualified_name), ) diff --git a/pyatlan/model/assets/__init__.py b/pyatlan/model/assets/__init__.py index 990d2d38e..54d13cfe2 100644 --- a/pyatlan/model/assets/__init__.py +++ b/pyatlan/model/assets/__init__.py @@ -66,8 +66,8 @@ "Column", "DatabricksUnityCatalogTag", "SnowflakeStream", - "Database", "CalculationView", + "Database", "Procedure", "SnowflakeTag", "CosmosMongoDB", @@ -161,6 +161,7 @@ "cognos": ["Cognos"], "superset": ["Superset"], "qlik": ["Qlik"], + "dataverse": ["Dataverse"], "cognite": ["Cognite"], "salesforce": ["Salesforce"], "readme_template": ["ReadmeTemplate"], @@ -285,6 +286,8 @@ "qlik_chart": ["QlikChart"], "qlik_dataset": ["QlikDataset"], "qlik_sheet": ["QlikSheet"], + "dataverse_attribute": ["DataverseAttribute"], + "dataverse_entity": ["DataverseEntity"], "cognite_event": ["CogniteEvent"], "cognite_asset": ["CogniteAsset"], "cognite_sequence": ["CogniteSequence"], diff --git a/pyatlan/model/assets/__init__.pyi b/pyatlan/model/assets/__init__.pyi index 6de750643..686622375 100644 --- a/pyatlan/model/assets/__init__.pyi +++ b/pyatlan/model/assets/__init__.pyi @@ -63,8 +63,8 @@ __all__ = [ "Column", "DatabricksUnityCatalogTag", "SnowflakeStream", - "Database", "CalculationView", + "Database", "Procedure", "SnowflakeTag", "CosmosMongoDB", @@ -158,6 +158,7 @@ __all__ = [ "Cognos", "Superset", "Qlik", + "Dataverse", "Cognite", "Salesforce", "ReadmeTemplate", @@ -282,6 +283,8 @@ __all__ = [ "QlikChart", "QlikDataset", "QlikSheet", + "DataverseAttribute", + "DataverseEntity", "CogniteEvent", "CogniteAsset", "CogniteSequence", @@ -474,6 +477,9 @@ from .cube_hierarchy import CubeHierarchy from .data_set import DataSet from .data_studio import DataStudio from .data_studio_asset import DataStudioAsset +from .dataverse import Dataverse +from .dataverse_attribute import DataverseAttribute +from .dataverse_entity import DataverseEntity from .dbt_column_process import DbtColumnProcess from .dbt_process import DbtProcess from .dbt_tag import DbtTag diff --git a/pyatlan/model/assets/core/__init__.py b/pyatlan/model/assets/core/__init__.py index 05b2fdd32..933105b6d 100644 --- a/pyatlan/model/assets/core/__init__.py +++ b/pyatlan/model/assets/core/__init__.py @@ -173,8 +173,8 @@ Column.Attributes.update_forward_refs(**localns) DatabricksUnityCatalogTag.Attributes.update_forward_refs(**localns) SnowflakeStream.Attributes.update_forward_refs(**localns) -Database.Attributes.update_forward_refs(**localns) CalculationView.Attributes.update_forward_refs(**localns) +Database.Attributes.update_forward_refs(**localns) Procedure.Attributes.update_forward_refs(**localns) SnowflakeTag.Attributes.update_forward_refs(**localns) CosmosMongoDB.Attributes.update_forward_refs(**localns) diff --git a/pyatlan/model/assets/core/procedure.py b/pyatlan/model/assets/core/procedure.py index 95c571edc..377491865 100644 --- a/pyatlan/model/assets/core/procedure.py +++ b/pyatlan/model/assets/core/procedure.py @@ -8,10 +8,10 @@ from pydantic.v1 import Field, validator +from pyatlan.model.enums import AtlanConnectorType from pyatlan.model.fields.atlan_fields import RelationField, TextField from pyatlan.utils import init_guid, validate_required_fields -from ...enums import AtlanConnectorType from .s_q_l import SQL diff --git a/pyatlan/model/structs.py b/pyatlan/model/structs.py index f72bdd76b..42df88751 100644 --- a/pyatlan/model/structs.py +++ b/pyatlan/model/structs.py @@ -42,19 +42,6 @@ def flatten_structs_attributes(cls, values: Dict[str, Any]) -> Dict[str, Any]: return values -class MCRuleSchedule(AtlanObject): - """Description""" - - mc_rule_schedule_type: Optional[str] = Field(default=None, description="") - mc_rule_schedule_interval_in_minutes: Optional[int] = Field( - default=None, description="" - ) - mc_rule_schedule_start_time: Optional[datetime] = Field( - default=None, description="" - ) - mc_rule_schedule_crontab: Optional[str] = Field(default=None, description="") - - class DbtJobRun(AtlanObject): """Description""" @@ -71,6 +58,19 @@ class DbtJobRun(AtlanObject): dbt_compiled_code: Optional[str] = Field(default=None, description="") +class MCRuleSchedule(AtlanObject): + """Description""" + + mc_rule_schedule_type: Optional[str] = Field(default=None, description="") + mc_rule_schedule_interval_in_minutes: Optional[int] = Field( + default=None, description="" + ) + mc_rule_schedule_start_time: Optional[datetime] = Field( + default=None, description="" + ) + mc_rule_schedule_crontab: Optional[str] = Field(default=None, description="") + + class AwsCloudWatchMetric(AtlanObject): """Description""" @@ -87,13 +87,6 @@ class KafkaTopicConsumption(AtlanObject): topic_current_offset: Optional[int] = Field(default=None, description="") -class Histogram(AtlanObject): - """Description""" - - boundaries: Set[float] = Field(description="") - frequencies: Set[float] = Field(description="") - - class Action(AtlanObject): """Description""" @@ -103,6 +96,13 @@ class Action(AtlanObject): task_action_display_text: Optional[str] = Field(default=None, description="") +class Histogram(AtlanObject): + """Description""" + + boundaries: Set[float] = Field(description="") + frequencies: Set[float] = Field(description="") + + class ColumnValueFrequencyMap(AtlanObject): """Description""" @@ -117,40 +117,6 @@ class SourceTagAttachmentValue(AtlanObject): tag_attachment_value: Optional[str] = Field(default=None, description="") -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 SourceTagAttachment(AtlanObject): """Description""" @@ -295,6 +261,40 @@ 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""" @@ -396,26 +396,26 @@ class SourceTagAttribute(AtlanObject): ) -MCRuleSchedule.update_forward_refs() - DbtJobRun.update_forward_refs() +MCRuleSchedule.update_forward_refs() + AwsCloudWatchMetric.update_forward_refs() KafkaTopicConsumption.update_forward_refs() -Histogram.update_forward_refs() - Action.update_forward_refs() +Histogram.update_forward_refs() + ColumnValueFrequencyMap.update_forward_refs() SourceTagAttachmentValue.update_forward_refs() -BadgeCondition.update_forward_refs() - SourceTagAttachment.update_forward_refs() +BadgeCondition.update_forward_refs() + AzureTag.update_forward_refs() StarredDetails.update_forward_refs() From b9004f31d474c00f5494bad6356ec2d6091485e5 Mon Sep 17 00:00:00 2001 From: Ernest Hill Date: Tue, 7 Jan 2025 15:30:42 +0300 Subject: [PATCH 4/6] Add missing classes --- pyatlan/model/assets/dataverse.py | 110 +++++++++++ pyatlan/model/assets/dataverse_attribute.py | 203 ++++++++++++++++++++ pyatlan/model/assets/dataverse_entity.py | 117 +++++++++++ 3 files changed, 430 insertions(+) create mode 100644 pyatlan/model/assets/dataverse.py create mode 100644 pyatlan/model/assets/dataverse_attribute.py create mode 100644 pyatlan/model/assets/dataverse_entity.py diff --git a/pyatlan/model/assets/dataverse.py b/pyatlan/model/assets/dataverse.py new file mode 100644 index 000000000..165651951 --- /dev/null +++ b/pyatlan/model/assets/dataverse.py @@ -0,0 +1,110 @@ +# 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 BooleanField + +from .saa_s import SaaS + + +class Dataverse(SaaS): + """Description""" + + type_name: str = Field(default="Dataverse", allow_mutation=False) + + @validator("type_name") + def validate_type_name(cls, v): + if v != "Dataverse": + raise ValueError("must be Dataverse") + return v + + def __setattr__(self, name, value): + if name in Dataverse._convenience_properties: + return object.__setattr__(self, name, value) + super().__setattr__(name, value) + + DATAVERSE_IS_CUSTOM: ClassVar[BooleanField] = BooleanField( + "dataverseIsCustom", "dataverseIsCustom" + ) + """ + Indicator if DataverseEntity is custom built. + """ + DATAVERSE_IS_CUSTOMIZABLE: ClassVar[BooleanField] = BooleanField( + "dataverseIsCustomizable", "dataverseIsCustomizable" + ) + """ + Indicator if DataverseEntity is customizable. + """ + DATAVERSE_IS_AUDIT_ENABLED: ClassVar[BooleanField] = BooleanField( + "dataverseIsAuditEnabled", "dataverseIsAuditEnabled" + ) + """ + Indicator if DataverseEntity has auditing enabled. + """ + + _convenience_properties: ClassVar[List[str]] = [ + "dataverse_is_custom", + "dataverse_is_customizable", + "dataverse_is_audit_enabled", + ] + + @property + def dataverse_is_custom(self) -> Optional[bool]: + return None if self.attributes is None else self.attributes.dataverse_is_custom + + @dataverse_is_custom.setter + def dataverse_is_custom(self, dataverse_is_custom: Optional[bool]): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_is_custom = dataverse_is_custom + + @property + def dataverse_is_customizable(self) -> Optional[bool]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_is_customizable + ) + + @dataverse_is_customizable.setter + def dataverse_is_customizable(self, dataverse_is_customizable: Optional[bool]): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_is_customizable = dataverse_is_customizable + + @property + def dataverse_is_audit_enabled(self) -> Optional[bool]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_is_audit_enabled + ) + + @dataverse_is_audit_enabled.setter + def dataverse_is_audit_enabled(self, dataverse_is_audit_enabled: Optional[bool]): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_is_audit_enabled = dataverse_is_audit_enabled + + class Attributes(SaaS.Attributes): + dataverse_is_custom: Optional[bool] = Field(default=None, description="") + dataverse_is_customizable: Optional[bool] = Field(default=None, description="") + dataverse_is_audit_enabled: Optional[bool] = Field(default=None, description="") + + attributes: Dataverse.Attributes = Field( + default_factory=lambda: Dataverse.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." + ), + ) + + +Dataverse.Attributes.update_forward_refs() diff --git a/pyatlan/model/assets/dataverse_attribute.py b/pyatlan/model/assets/dataverse_attribute.py new file mode 100644 index 000000000..1c74130ff --- /dev/null +++ b/pyatlan/model/assets/dataverse_attribute.py @@ -0,0 +1,203 @@ +# 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 BooleanField, KeywordField, RelationField + +from .dataverse import Dataverse + + +class DataverseAttribute(Dataverse): + """Description""" + + type_name: str = Field(default="DataverseAttribute", allow_mutation=False) + + @validator("type_name") + def validate_type_name(cls, v): + if v != "DataverseAttribute": + raise ValueError("must be DataverseAttribute") + return v + + def __setattr__(self, name, value): + if name in DataverseAttribute._convenience_properties: + return object.__setattr__(self, name, value) + super().__setattr__(name, value) + + DATAVERSE_ENTITY_QUALIFIED_NAME: ClassVar[KeywordField] = KeywordField( + "dataverseEntityQualifiedName", "dataverseEntityQualifiedName" + ) + """ + Entity Qualified Name of the DataverseAttribute. + """ + DATAVERSE_ATTRIBUTE_SCHEMA_NAME: ClassVar[KeywordField] = KeywordField( + "dataverseAttributeSchemaName", "dataverseAttributeSchemaName" + ) + """ + Schema Name of the DataverseAttribute. + """ + DATAVERSE_ATTRIBUTE_TYPE: ClassVar[KeywordField] = KeywordField( + "dataverseAttributeType", "dataverseAttributeType" + ) + """ + Type of the DataverseAttribute. + """ + DATAVERSE_ATTRIBUTE_IS_PRIMARY_ID: ClassVar[BooleanField] = BooleanField( + "dataverseAttributeIsPrimaryId", "dataverseAttributeIsPrimaryId" + ) + """ + Indicator if DataverseAttribute is the primary key. + """ + DATAVERSE_ATTRIBUTE_IS_SEARCHABLE: ClassVar[BooleanField] = BooleanField( + "dataverseAttributeIsSearchable", "dataverseAttributeIsSearchable" + ) + """ + Indicator if DataverseAttribute is searchable. + """ + + DATAVERSE_ENTITY: ClassVar[RelationField] = RelationField("dataverseEntity") + """ + TBC + """ + + _convenience_properties: ClassVar[List[str]] = [ + "dataverse_entity_qualified_name", + "dataverse_attribute_schema_name", + "dataverse_attribute_type", + "dataverse_attribute_is_primary_id", + "dataverse_attribute_is_searchable", + "dataverse_entity", + ] + + @property + def dataverse_entity_qualified_name(self) -> Optional[str]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_entity_qualified_name + ) + + @dataverse_entity_qualified_name.setter + def dataverse_entity_qualified_name( + self, dataverse_entity_qualified_name: Optional[str] + ): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_entity_qualified_name = ( + dataverse_entity_qualified_name + ) + + @property + def dataverse_attribute_schema_name(self) -> Optional[str]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_attribute_schema_name + ) + + @dataverse_attribute_schema_name.setter + def dataverse_attribute_schema_name( + self, dataverse_attribute_schema_name: Optional[str] + ): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_attribute_schema_name = ( + dataverse_attribute_schema_name + ) + + @property + def dataverse_attribute_type(self) -> Optional[str]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_attribute_type + ) + + @dataverse_attribute_type.setter + def dataverse_attribute_type(self, dataverse_attribute_type: Optional[str]): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_attribute_type = dataverse_attribute_type + + @property + def dataverse_attribute_is_primary_id(self) -> Optional[bool]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_attribute_is_primary_id + ) + + @dataverse_attribute_is_primary_id.setter + def dataverse_attribute_is_primary_id( + self, dataverse_attribute_is_primary_id: Optional[bool] + ): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_attribute_is_primary_id = ( + dataverse_attribute_is_primary_id + ) + + @property + def dataverse_attribute_is_searchable(self) -> Optional[bool]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_attribute_is_searchable + ) + + @dataverse_attribute_is_searchable.setter + def dataverse_attribute_is_searchable( + self, dataverse_attribute_is_searchable: Optional[bool] + ): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_attribute_is_searchable = ( + dataverse_attribute_is_searchable + ) + + @property + def dataverse_entity(self) -> Optional[DataverseEntity]: + return None if self.attributes is None else self.attributes.dataverse_entity + + @dataverse_entity.setter + def dataverse_entity(self, dataverse_entity: Optional[DataverseEntity]): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_entity = dataverse_entity + + class Attributes(Dataverse.Attributes): + dataverse_entity_qualified_name: Optional[str] = Field( + default=None, description="" + ) + dataverse_attribute_schema_name: Optional[str] = Field( + default=None, description="" + ) + dataverse_attribute_type: Optional[str] = Field(default=None, description="") + dataverse_attribute_is_primary_id: Optional[bool] = Field( + default=None, description="" + ) + dataverse_attribute_is_searchable: Optional[bool] = Field( + default=None, description="" + ) + dataverse_entity: Optional[DataverseEntity] = Field( + default=None, description="" + ) # relationship + + attributes: DataverseAttribute.Attributes = Field( + default_factory=lambda: DataverseAttribute.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 .dataverse_entity import DataverseEntity # noqa + +DataverseAttribute.Attributes.update_forward_refs() diff --git a/pyatlan/model/assets/dataverse_entity.py b/pyatlan/model/assets/dataverse_entity.py new file mode 100644 index 000000000..ecbd9ede0 --- /dev/null +++ b/pyatlan/model/assets/dataverse_entity.py @@ -0,0 +1,117 @@ +# 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 + +from .dataverse import Dataverse + + +class DataverseEntity(Dataverse): + """Description""" + + type_name: str = Field(default="DataverseEntity", allow_mutation=False) + + @validator("type_name") + def validate_type_name(cls, v): + if v != "DataverseEntity": + raise ValueError("must be DataverseEntity") + return v + + def __setattr__(self, name, value): + if name in DataverseEntity._convenience_properties: + return object.__setattr__(self, name, value) + super().__setattr__(name, value) + + DATAVERSE_ENTITY_SCHEMA_NAME: ClassVar[KeywordField] = KeywordField( + "dataverseEntitySchemaName", "dataverseEntitySchemaName" + ) + """ + Schema Name of the DataverseEntity. + """ + DATAVERSE_ENTITY_TABLE_TYPE: ClassVar[KeywordField] = KeywordField( + "dataverseEntityTableType", "dataverseEntityTableType" + ) + """ + Table Type of the DataverseEntity. + """ + + DATAVERSE_ATTRIBUTES: ClassVar[RelationField] = RelationField("dataverseAttributes") + """ + TBC + """ + + _convenience_properties: ClassVar[List[str]] = [ + "dataverse_entity_schema_name", + "dataverse_entity_table_type", + "dataverse_attributes", + ] + + @property + def dataverse_entity_schema_name(self) -> Optional[str]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_entity_schema_name + ) + + @dataverse_entity_schema_name.setter + def dataverse_entity_schema_name(self, dataverse_entity_schema_name: Optional[str]): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_entity_schema_name = dataverse_entity_schema_name + + @property + def dataverse_entity_table_type(self) -> Optional[str]: + return ( + None + if self.attributes is None + else self.attributes.dataverse_entity_table_type + ) + + @dataverse_entity_table_type.setter + def dataverse_entity_table_type(self, dataverse_entity_table_type: Optional[str]): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_entity_table_type = dataverse_entity_table_type + + @property + def dataverse_attributes(self) -> Optional[List[DataverseAttribute]]: + return None if self.attributes is None else self.attributes.dataverse_attributes + + @dataverse_attributes.setter + def dataverse_attributes( + self, dataverse_attributes: Optional[List[DataverseAttribute]] + ): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.dataverse_attributes = dataverse_attributes + + class Attributes(Dataverse.Attributes): + dataverse_entity_schema_name: Optional[str] = Field( + default=None, description="" + ) + dataverse_entity_table_type: Optional[str] = Field(default=None, description="") + dataverse_attributes: Optional[List[DataverseAttribute]] = Field( + default=None, description="" + ) # relationship + + attributes: DataverseEntity.Attributes = Field( + default_factory=lambda: DataverseEntity.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 .dataverse_attribute import DataverseAttribute # noqa + +DataverseEntity.Attributes.update_forward_refs() From 56ba4b5d5f0d0fdcebb94668f5b0b7a7202879cc Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Tue, 7 Jan 2025 19:24:03 +0530 Subject: [PATCH 5/6] [tests] Added integration tests for `Procedure.creator/updater()` methods --- .../templates/methods/asset/procedure.jinja2 | 20 ++++ pyatlan/model/assets/core/procedure.py | 20 ++++ tests/integration/test_sql_assets.py | 93 +++++++++++++++++++ 3 files changed, 133 insertions(+) diff --git a/pyatlan/generator/templates/methods/asset/procedure.jinja2 b/pyatlan/generator/templates/methods/asset/procedure.jinja2 index f4dcae109..ff4a14e29 100644 --- a/pyatlan/generator/templates/methods/asset/procedure.jinja2 +++ b/pyatlan/generator/templates/methods/asset/procedure.jinja2 @@ -46,3 +46,23 @@ connection_qualified_name=connection_qualified_name, ) return cls(attributes=attributes) + + @classmethod + @init_guid + def updater(cls, *, name: str, qualified_name: str, definition: str) -> Procedure: + validate_required_fields( + ["name", "qualified_name", "definition"], + [name, qualified_name, definition], + ) + procedure = Procedure( + attributes=Procedure.Attributes(qualified_name=qualified_name, name=name) + ) + procedure.definition = definition + return procedure + + def trim_to_required(self: Procedure) -> Procedure: + return self.updater( + qualified_name=self.qualified_name or "", + name=self.name or "", + definition=self.definition or "", + ) diff --git a/pyatlan/model/assets/core/procedure.py b/pyatlan/model/assets/core/procedure.py index 377491865..52e9f238f 100644 --- a/pyatlan/model/assets/core/procedure.py +++ b/pyatlan/model/assets/core/procedure.py @@ -66,6 +66,26 @@ def creator( ) return cls(attributes=attributes) + @classmethod + @init_guid + def updater(cls, *, name: str, qualified_name: str, definition: str) -> Procedure: + validate_required_fields( + ["name", "qualified_name", "definition"], + [name, qualified_name, definition], + ) + procedure = Procedure( + attributes=Procedure.Attributes(qualified_name=qualified_name, name=name) + ) + procedure.definition = definition + return procedure + + def trim_to_required(self: Procedure) -> Procedure: + return self.updater( + qualified_name=self.qualified_name or "", + name=self.name or "", + definition=self.definition or "", + ) + type_name: str = Field(default="Procedure", allow_mutation=False) @validator("type_name") diff --git a/tests/integration/test_sql_assets.py b/tests/integration/test_sql_assets.py index 1d2a90c70..c3e9b09e1 100644 --- a/tests/integration/test_sql_assets.py +++ b/tests/integration/test_sql_assets.py @@ -12,6 +12,7 @@ Column, Connection, Database, + Procedure, Readme, Schema, Table, @@ -548,6 +549,98 @@ def test_trim_to_required( assert response.mutated_entities is None +@pytest.mark.order(after="TestView") +class TestProcedure: + procedure: Optional[Procedure] = None + _DEFINITION = """ + BEGIN + insert into `atlanhq.testing_lineage.INSTACART_ALCOHOL_ORDER_TIME_copy` + select * from `atlanhq.testing_lineage.INSTACART_ALCOHOL_ORDER_TIME`; + END + """ + + def test_creator( + self, + client: AtlanClient, + upsert: Callable[[Asset], AssetMutationResponse], + ): + procedure_name = TestId.make_unique("My_Procedure") + assert TestSchema.schema is not None + assert TestSchema.schema.qualified_name + procedure = Procedure.creator( + name=procedure_name, + definition=self._DEFINITION, + schema_qualified_name=TestSchema.schema.qualified_name, + ) + response = upsert(procedure) + assert response.mutated_entities + assert response.mutated_entities.CREATE + assert len(response.mutated_entities.CREATE) == 1 + assert isinstance(response.mutated_entities.CREATE[0], Procedure) + assert response.guid_assignments + procedure = response.mutated_entities.CREATE[0] + TestProcedure.procedure = procedure + + def test_overload_creator( + self, + client: AtlanClient, + upsert: Callable[[Asset], AssetMutationResponse], + ): + procedure_name = TestId.make_unique("My_Procedure_Overload") + assert TestDatabase.database is not None + assert TestDatabase.database.name + assert TestDatabase.database.qualified_name + assert TestSchema.schema is not None + assert TestSchema.schema.name + assert TestSchema.schema.qualified_name + assert TestConnection.connection is not None + assert TestConnection.connection.qualified_name + + procedure = Procedure.creator( + name=procedure_name, + definition=self._DEFINITION, + schema_name=TestSchema.schema.name, + schema_qualified_name=TestSchema.schema.qualified_name, + database_name=TestDatabase.database.name, + database_qualified_name=TestDatabase.database.qualified_name, + connection_qualified_name=TestConnection.connection.qualified_name, + ) + response = upsert(procedure) + assert response.mutated_entities + assert response.mutated_entities.CREATE + assert len(response.mutated_entities.CREATE) == 1 + assert isinstance(response.mutated_entities.CREATE[0], Procedure) + assert response.guid_assignments + + @pytest.mark.order(after="test_creator") + def test_updater( + self, client: AtlanClient, upsert: Callable[[Asset], AssetMutationResponse] + ): + assert TestProcedure.procedure + procedure = TestProcedure.procedure + assert procedure.qualified_name + assert procedure.name + assert procedure.definition + description = f"{procedure.description} more stuff" + procedure = Procedure.updater( + qualified_name=procedure.qualified_name, + name=procedure.name, + definition=procedure.definition, + ) + procedure.description = description + response = upsert(procedure) + verify_asset_updated(response, Procedure) + + @pytest.mark.order(after="test_creator") + def test_trim_to_required( + self, client: AtlanClient, upsert: Callable[[Asset], AssetMutationResponse] + ): + assert TestProcedure.procedure + procedure = TestProcedure.procedure.trim_to_required() + response = upsert(procedure) + assert response.mutated_entities is None + + @pytest.mark.order(after="TestView") class TestColumn: column: Optional[Column] = None From 336ce7caa61f3ddc323e10bcc0bd07363515739f Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Tue, 7 Jan 2025 19:26:59 +0530 Subject: [PATCH 6/6] [gen] Re-ran generator with the latest typedefs --- docs/asset/dataverse.rst | 10 ++++++++++ docs/asset/dataverseattribute.rst | 10 ++++++++++ docs/asset/dataverseentity.rst | 10 ++++++++++ 3 files changed, 30 insertions(+) create mode 100644 docs/asset/dataverse.rst create mode 100644 docs/asset/dataverseattribute.rst create mode 100644 docs/asset/dataverseentity.rst diff --git a/docs/asset/dataverse.rst b/docs/asset/dataverse.rst new file mode 100644 index 000000000..498dd79d8 --- /dev/null +++ b/docs/asset/dataverse.rst @@ -0,0 +1,10 @@ +.. _dataverse: + +Dataverse +========= + +.. module:: pyatlan.model.assets + :no-index: + +.. autoclass:: Dataverse + :members: diff --git a/docs/asset/dataverseattribute.rst b/docs/asset/dataverseattribute.rst new file mode 100644 index 000000000..f6923476c --- /dev/null +++ b/docs/asset/dataverseattribute.rst @@ -0,0 +1,10 @@ +.. _dataverseattribute: + +DataverseAttribute +================== + +.. module:: pyatlan.model.assets + :no-index: + +.. autoclass:: DataverseAttribute + :members: diff --git a/docs/asset/dataverseentity.rst b/docs/asset/dataverseentity.rst new file mode 100644 index 000000000..d18cc1c72 --- /dev/null +++ b/docs/asset/dataverseentity.rst @@ -0,0 +1,10 @@ +.. _dataverseentity: + +DataverseEntity +=============== + +.. module:: pyatlan.model.assets + :no-index: + +.. autoclass:: DataverseEntity + :members: