Skip to content

Commit b47a712

Browse files
committed
[feat] Migrate v9 models and tests: events, packages, open lineage, DQ rules, lineage
- pyatlan_v9/model/events.py: Full v9 migration with AtlanEvent.from_dict() for polymorphic Asset dispatch via type registry and payload discrimination - pyatlan_v9/model/packages/: Migrate AbstractPackage, crawlers, miners to msgspec - pyatlan_v9/model/open_lineage/: All OpenLineage models as msgspec.Struct - pyatlan_v9/model/assets/data_quality_rule.py: Creator/updater methods, static helpers - pyatlan_v9/model/assets/asset.py: remove_description/user_description/owners, ClassVar descriptors - pyatlan_v9/model/workflow.py: rename=camel, to_json() - pyatlan_v9/model/credential.py: rename=camel - pyatlan_v9/model/lineage.py: rename=camel, validate_arguments - pyatlan_v9/client/atlan.py: msgspec.Struct handling in _create_params, parse_query, upload_image - tests_v9/unit/: Update all v9 tests to use v9 models and client exclusively - tests_v9/unit/test_events.py: Uses v9 AtlanEvent.from_dict(), no legacy imports - tests_v9/unit/test_lineage.py, test_model.py: New, ported from legacy
1 parent 5859f89 commit b47a712

70 files changed

Lines changed: 9258 additions & 387 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

pyatlan_v9/client/atlan.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,8 @@
5353
from pyatlan.client.user import UserClient
5454
from pyatlan.client.workflow import WorkflowClient
5555
from pyatlan.errors import ERROR_CODE_FOR_HTTP_STATUS, AtlanError, ErrorCode
56-
from pyatlan.model.atlan_image import AtlanImage
5756
from pyatlan.model.core import AtlanObject, AtlanRequest, AtlanResponse
5857
from pyatlan.model.enums import AtlanTypeCategory
59-
from pyatlan.model.query import ParsedQuery, QueryParserRequest
6058
from pyatlan.multipart_data_generator import MultipartDataGenerator
6159
from pyatlan.utils import (
6260
API,
@@ -66,6 +64,8 @@
6664
get_python_version,
6765
)
6866
from pyatlan_v9.client.transport import PyatlanSyncTransport
67+
from pyatlan_v9.model.atlan_image import AtlanImage
68+
from pyatlan_v9.model.query import ParsedQuery, QueryParserRequest
6969

7070
request_id_var = ContextVar("request_id", default=None)
7171

@@ -657,6 +657,8 @@ def _create_params(
657657
params["data"] = AtlanRequest(instance=request_obj, client=self).json()
658658
elif api.consumes == APPLICATION_ENCODED_FORM:
659659
params["data"] = request_obj
660+
elif isinstance(request_obj, msgspec.Struct):
661+
params["data"] = msgspec.json.encode(request_obj)
660662
else:
661663
params["data"] = json.dumps(request_obj)
662664
return params
@@ -736,7 +738,7 @@ def upload_image(self, file, filename: str) -> AtlanImage:
736738
:raises AtlanError: on any API communication issue
737739
"""
738740
raw_json = self._upload_file(UPLOAD_IMAGE, file=file, filename=filename)
739-
return AtlanImage(**raw_json)
741+
return msgspec.convert(raw_json, AtlanImage, strict=False)
740742

741743
def parse_query(self, query: QueryParserRequest) -> Optional[ParsedQuery]:
742744
"""
@@ -751,7 +753,7 @@ def parse_query(self, query: QueryParserRequest) -> Optional[ParsedQuery]:
751753
request_obj=query,
752754
exclude_unset=True,
753755
)
754-
return ParsedQuery(**raw_json)
756+
return msgspec.convert(raw_json, ParsedQuery, strict=False)
755757

756758
@contextlib.contextmanager
757759
def max_retries(

pyatlan_v9/model/assets/asset.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,9 @@ class Asset(Referenceable):
5050
DISPLAY_NAME: ClassVar = None # Initialized after imports below
5151
CONNECTOR_NAME: ClassVar = None # Initialized after imports below
5252
CONNECTION_NAME: ClassVar = None # Initialized after imports below
53+
ASSET_DQ_ROW_SCOPE_FILTER_COLUMN_QUALIFIED_NAME: ClassVar = (
54+
None # Initialized after imports below
55+
)
5356

5457
# =========================================================================
5558
# Instance Fields
@@ -743,6 +746,37 @@ def remove_certificate(self) -> "Asset":
743746
self.certificate_status_message = None
744747
return self
745748

749+
def remove_description(self) -> "Asset":
750+
"""
751+
Remove the description from this asset.
752+
753+
Returns:
754+
Self for fluent chaining
755+
"""
756+
self.description = None
757+
return self
758+
759+
def remove_user_description(self) -> "Asset":
760+
"""
761+
Remove the user description from this asset.
762+
763+
Returns:
764+
Self for fluent chaining
765+
"""
766+
self.user_description = None
767+
return self
768+
769+
def remove_owners(self) -> "Asset":
770+
"""
771+
Remove the owners from this asset.
772+
773+
Returns:
774+
Self for fluent chaining
775+
"""
776+
self.owner_groups = None
777+
self.owner_users = None
778+
return self
779+
746780
def flush_custom_metadata(self, client=None) -> None:
747781
"""
748782
Flush (clear) all custom metadata on this asset.
@@ -816,6 +850,8 @@ def trim_to_required(self) -> "Asset":
816850
KeywordField,
817851
KeywordTextField,
818852
KeywordTextStemmedField,
853+
NumericField,
854+
TextField,
819855
)
820856

821857
Asset.NAME = KeywordTextStemmedField("name", "name.keyword", "name", "name.stemmed")
@@ -826,6 +862,32 @@ def trim_to_required(self) -> "Asset":
826862
Asset.CONNECTION_NAME = KeywordTextField(
827863
"connectionName", "connectionName", "connectionName.text"
828864
)
865+
Asset.DESCRIPTION = KeywordTextField(
866+
"description", "description.keyword", "description.text"
867+
)
868+
Asset.USER_DESCRIPTION = KeywordTextField(
869+
"userDescription", "userDescription.keyword", "userDescription.text"
870+
)
871+
Asset.CERTIFICATE_STATUS = KeywordTextField(
872+
"certificateStatus", "certificateStatus", "certificateStatus.text"
873+
)
874+
Asset.CERTIFICATE_STATUS_MESSAGE = TextField(
875+
"certificateStatusMessage", "certificateStatusMessage"
876+
)
877+
Asset.ANNOUNCEMENT_TITLE = TextField("announcementTitle", "announcementTitle")
878+
Asset.ANNOUNCEMENT_MESSAGE = TextField("announcementMessage", "announcementMessage")
879+
Asset.ANNOUNCEMENT_TYPE = KeywordField("announcementType", "announcementType")
880+
Asset.ANNOUNCEMENT_UPDATED_AT = NumericField(
881+
"announcementUpdatedAt", "announcementUpdatedAt"
882+
)
883+
Asset.OWNER_USERS = KeywordField("ownerUsers", "ownerUsers")
884+
Asset.OWNER_GROUPS = KeywordField("ownerGroups", "ownerGroups")
885+
Asset.ADMIN_USERS = KeywordField("adminUsers", "adminUsers")
886+
Asset.ADMIN_GROUPS = KeywordField("adminGroups", "adminGroups")
887+
Asset.ASSET_DQ_ROW_SCOPE_FILTER_COLUMN_QUALIFIED_NAME = KeywordField(
888+
"assetDQRowScopeFilterColumnQualifiedName",
889+
"assetDQRowScopeFilterColumnQualifiedName",
890+
)
829891

830892

831893
# =============================================================================

pyatlan_v9/model/assets/connection.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -206,17 +206,20 @@ def creator(
206206
client.role_cache.validate_idstrs(idstrs=admin_roles or [])
207207
client.group_cache.validate_aliases(aliases=admin_groups or [])
208208

209-
return cls(
209+
kwargs: dict = dict(
210210
name=name,
211211
qualified_name=connector_type.to_qualified_name(),
212212
connector_name=connector_type.value,
213213
category=connector_type.category.value,
214214
admin_users=set() if admin_users is None else set(admin_users),
215215
admin_groups=set() if admin_groups is None else set(admin_groups),
216216
admin_roles=set() if admin_roles is None else set(admin_roles),
217-
host=host,
218-
port=port,
219217
)
218+
if host is not None:
219+
kwargs["host"] = host
220+
if port is not None:
221+
kwargs["port"] = port
222+
return cls(**kwargs)
220223

221224
@classmethod
222225
def updater(cls, *, qualified_name: str, name: str) -> "Connection":

0 commit comments

Comments
 (0)