Merge branch 'main' into APP-5632

atlanhq · Feb 25, 2025 · 7a6bc66 · 7a6bc66
2 parents e9198c9 + 3ee3477
commit 7a6bc66
Show file tree

Hide file tree

Showing 397 changed files with 4,205 additions and 1,494 deletions.
diff --git a/.github/workflows/pyatlan-pr.yaml b/.github/workflows/pyatlan-pr.yaml
@@ -60,16 +60,14 @@ jobs:
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
           if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
 
-      - name: QA checks (black, flake8, mypy)
+      - name: QA checks (ruff-format, ruff-lint, mypy)
         run: |
           ./qa-checks
 
       - name: Run unit tests
         env: # Test tenant environment variables
           ATLAN_API_KEY: ${{ secrets.ATLAN_API_KEY }}
           ATLAN_BASE_URL: ${{ secrets.ATLAN_BASE_URL }}
-          MARK_API_KEY: ${{ secrets.MARK_ATLAN_API_KEY }}
-          MARK_BASE_URL: https://mark.atlan.com
         # Run with `pytest-sugar` for enhancing the overall test report output
         run: pytest tests/unit --force-sugar
 

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,42 +1,17 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: check-yaml
       - id: end-of-file-fixer
       - id: trailing-whitespace
       - id: debug-statements
 
-  - repo: https://github.com/PyCQA/autoflake
-    rev: v2.3.1
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.9.7
     hooks:
-    -   id: autoflake
-        args: [--ignore-init-module-imports, --remove-all-unused-imports, --remove-unused-variables, --in-place]
-        language_version: python3
-
-  - repo: https://github.com/pycqa/flake8
-    rev: 7.0.0
-    hooks:
-      - id: flake8
-        additional_dependencies: ["flake8-bandit", "flake8-bugbear"]
-
-  - repo: https://github.com/pycqa/isort
-    rev: 5.13.2
-    hooks:
-      - id: isort
-        args: ["--profile", "black", "--filter-files"]
-
-  - repo: https://github.com/psf/black
-    rev: 24.4.2
-    hooks:
-      - id: black
-        language_version: python3
-
-  - repo: local
-    hooks:
-      - id: tests
-        name: run tests
-        entry: pytest
-        language: system
-        types: [python]
-        stages: [push]
+      - id: ruff
+      - id: ruff
+        name: ruff-check-autofix
+        args: ["check", "--select", "I", "--fix", "--silent"]
+      - id: ruff-format
diff --git a/docs/asset/ai.rst b/docs/asset/ai.rst
@@ -0,0 +1,10 @@
+.. _ai:
+
+AI
+==
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: AI
+    :members:
diff --git a/docs/asset/aiapplication.rst b/docs/asset/aiapplication.rst
@@ -0,0 +1,10 @@
+.. _aiapplication:
+
+AIApplication
+=============
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: AIApplication
+    :members:
diff --git a/docs/asset/aimodel.rst b/docs/asset/aimodel.rst
@@ -0,0 +1,10 @@
+.. _aimodel:
+
+AIModel
+=======
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: AIModel
+    :members:
diff --git a/docs/asset/documentdb.rst b/docs/asset/documentdb.rst
@@ -0,0 +1,10 @@
+.. _documentdb:
+
+DocumentDB
+==========
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: DocumentDB
+    :members:
diff --git a/docs/asset/documentdbcollection.rst b/docs/asset/documentdbcollection.rst
@@ -0,0 +1,10 @@
+.. _documentdbcollection:
+
+DocumentDBCollection
+====================
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: DocumentDBCollection
+    :members:
diff --git a/docs/asset/documentdbdatabase.rst b/docs/asset/documentdbdatabase.rst
@@ -0,0 +1,10 @@
+.. _documentdbdatabase:
+
+DocumentDBDatabase
+==================
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: DocumentDBDatabase
+    :members:
diff --git a/docs/asset/sourcetag.rst b/docs/asset/sourcetag.rst
@@ -1,10 +1,10 @@
 .. _sourcetag:
 
-sourceTag
+SourceTag
 =========
 
 .. module:: pyatlan.model.assets
     :no-index:
 
-.. autoclass:: sourceTag
+.. autoclass:: SourceTag
     :members:
diff --git a/docs/assets.rst b/docs/assets.rst
@@ -13,6 +13,9 @@ You can interact with all of the following different kinds of assets:
     asset/adlsaccount
     asset/adlscontainer
     asset/adlsobject
+    asset/ai
+    asset/aiapplication
+    asset/aimodel
     asset/api
     asset/apifield
     asset/apiobject
@@ -121,6 +124,9 @@ You can interact with all of the following different kinds of assets:
     asset/dbtsource
     asset/dbttag
     asset/dbttest
+    asset/documentdb
+    asset/documentdbcollection
+    asset/documentdbdatabase
     asset/domo
     asset/domocard
     asset/domodashboard
@@ -283,6 +289,7 @@ You can interact with all of the following different kinds of assets:
     asset/snowflaketag
     asset/soda
     asset/sodacheck
+    asset/sourcetag
     asset/spark
     asset/sparkjob
     asset/stakeholder
@@ -318,4 +325,3 @@ You can interact with all of the following different kinds of assets:
     asset/view
     asset/workflow
     asset/workflowrun
-    asset/sourcetag
diff --git a/pyatlan/cache/connection_cache.py b/pyatlan/cache/connection_cache.py
@@ -149,8 +149,7 @@ def lookup_by_name(self, name: ConnectionName) -> None:
         if len(results) > 1:
             LOGGER.warning(
                 (
-                    "Found multiple connections of the same type "
-                    "with the same name, caching only the first: %s"
+                    "Found multiple connections of the same type with the same name, caching only the first: %s"
                 ),
                 name,
             )

diff --git a/pyatlan/client/asset.py b/pyatlan/client/asset.py
@@ -506,7 +506,9 @@ def _handle_relationships(self, raw_json):
 
     @validate_arguments
     def retrieve_minimal(
-        self, guid: str, asset_type: Type[A] = Asset  # type: ignore[assignment]
+        self,
+        guid: str,
+        asset_type: Type[A] = Asset,  # type: ignore[assignment]
     ) -> A:
         """
         Retrieves an asset by its GUID, without any of its relationships.
@@ -599,8 +601,7 @@ def upsert_merging_cm(
     ) -> AssetMutationResponse:
         """Deprecated - use save_merging_cm() instead."""
         warn(
-            "This method is deprecated, please use 'save_merging_cm' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'save_merging_cm' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -658,8 +659,7 @@ def upsert_replacing_cm(
     ) -> AssetMutationResponse:
         """Deprecated - use save_replacing_cm() instead."""
         warn(
-            "This method is deprecated, please use 'save_replacing_cm' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'save_replacing_cm' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -1603,7 +1603,9 @@ def find_glossary_by_name(
     def find_category_fast_by_name(
         self,
         name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
-        glossary_qualified_name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
+        glossary_qualified_name: constr(  # type: ignore
+            strip_whitespace=True, min_length=1, strict=True
+        ),
         attributes: Optional[List[StrictStr]] = None,
     ) -> List[AtlasGlossaryCategory]:
         """
@@ -1699,7 +1701,9 @@ def _search_for_asset_with_name(
     def find_term_fast_by_name(
         self,
         name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
-        glossary_qualified_name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
+        glossary_qualified_name: constr(  # type: ignore
+            strip_whitespace=True, min_length=1, strict=True
+        ),
         attributes: Optional[List[StrictStr]] = None,
     ) -> AtlasGlossaryTerm:
         """

diff --git a/pyatlan/client/atlan.py b/pyatlan/client/atlan.py
@@ -647,8 +647,7 @@ def get_roles(
     def get_all_roles(self) -> RoleResponse:
         """Deprecated - use self.role.get_all() instead."""
         warn(
-            "This method is deprecated, please use 'self.role.get_all' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'self.role.get_all' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -783,8 +782,7 @@ def get_groups_for_user(
     ) -> GroupResponse:
         """Deprecated - use user.get_groups() instead."""
         warn(
-            "This method is deprecated, please use 'user.get_groups' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'user.get_groups' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -797,8 +795,7 @@ def add_user_to_groups(
     ) -> None:
         """Deprecated - use user.add_to_groups() instead."""
         warn(
-            "This method is deprecated, please use 'user.add_to_groups' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'user.add_to_groups' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -811,8 +808,7 @@ def change_user_role(
     ) -> None:
         """Deprecated - use user.change_role() instead."""
         warn(
-            "This method is deprecated, please use 'user.change_role' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'user.change_role' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -823,8 +819,7 @@ def get_current_user(
     ) -> UserMinimalResponse:
         """Deprecated - use user.get_current() instead."""
         warn(
-            "This method is deprecated, please use 'user.get_current' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'user.get_current' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -867,8 +862,7 @@ def get_users_by_email(
     ) -> Optional[List[AtlanUser]]:
         """Deprecated - use user.get_by_email() instead."""
         warn(
-            "This method is deprecated, please use 'user.get_by_email' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'user.get_by_email' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -932,8 +926,7 @@ def get_asset_by_guid(
     ) -> A:
         """Deprecated - use asset.get_by_guid() instead."""
         warn(
-            "This method is deprecated, please use 'asset.get_by_guid' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'asset.get_by_guid' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -1106,8 +1099,7 @@ def delete_entity_by_guid(
     def restore(self, asset_type: Type[A], qualified_name: str) -> bool:
         """Deprecated - use asset.restore() instead."""
         warn(
-            "This method is deprecated, please use 'asset.restore' instead, which offers identical "
-            "functionality.",
+            "This method is deprecated, please use 'asset.restore' instead, which offers identical functionality.",
             DeprecationWarning,
             stacklevel=2,
         )
@@ -1590,7 +1582,9 @@ def find_glossary_by_name(
     def find_category_fast_by_name(
         self,
         name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
-        glossary_qualified_name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
+        glossary_qualified_name: constr(  # type: ignore
+            strip_whitespace=True, min_length=1, strict=True
+        ),
         attributes: Optional[List[StrictStr]] = None,
     ) -> List[AtlasGlossaryCategory]:
         """Deprecated - use asset.find_category_fast_by_name() instead."""
@@ -1628,7 +1622,9 @@ def find_category_by_name(
     def find_term_fast_by_name(
         self,
         name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
-        glossary_qualified_name: constr(strip_whitespace=True, min_length=1, strict=True),  # type: ignore
+        glossary_qualified_name: constr(  # type: ignore
+            strip_whitespace=True, min_length=1, strict=True
+        ),
         attributes: Optional[List[StrictStr]] = None,
     ) -> AtlasGlossaryTerm:
         """Deprecated - use asset.find_category_by_name() instead."""
@@ -1672,8 +1668,8 @@ def max_retries(
             adapter = self._session.adapters[HTTP_PREFIX]
         else:
             adapter = self._session.adapters[HTTPS_PREFIX]
-        current_max = adapter.max_retries
-        adapter.max_retries = max_retries
+        current_max = adapter.max_retries  # type: ignore[attr-defined]
+        adapter.max_retries = max_retries  # type: ignore[attr-defined]
         LOGGER.debug(
             "max_retries set to total: %s force_list: %s",
             max_retries.total,
@@ -1687,11 +1683,11 @@ def max_retries(
             LOGGER.exception("Exception in max retries")
             raise ErrorCode.RETRY_OVERRUN.exception_with_parameters() from err
         finally:
-            adapter.max_retries = current_max
+            adapter.max_retries = current_max  # type: ignore[attr-defined]
             LOGGER.debug(
                 "max_retries restored to total: %s force_list: %s",
-                adapter.max_retries.total,
-                adapter.max_retries.status_forcelist,
+                adapter.max_retries.total,  # type: ignore[attr-defined]
+                adapter.max_retries.status_forcelist,  # type: ignore[attr-defined]
             )
 
 

diff --git a/pyatlan/client/constants.py b/pyatlan/client/constants.py
@@ -299,8 +299,8 @@
     endpoint=EndPoint.ATLAS,
 )
 DELETE_CLASSIFICATION_BY_TYPE_AND_ATTRIBUTE = API(
-    ENTITY_API + "uniqueAttribute/type/{type_name}/classification/{"
-    "classification_name}",
+    ENTITY_API
+    + "uniqueAttribute/type/{type_name}/classification/{classification_name}",
     HTTPMethod.DELETE,
     HTTPStatus.NO_CONTENT,
     endpoint=EndPoint.ATLAS,

diff --git a/pyatlan/errors.py b/pyatlan/errors.py
@@ -6,10 +6,7 @@
 from typing import Dict, List, Protocol, Type, TypeVar
 
 E = TypeVar("E", bound="AtlanError")
-RAISE_GITHUB_ISSUE = (
-    "Please raise an issue on the Python SDK GitHub "
-    "repository providing context in which this error occurred."
-)
+RAISE_GITHUB_ISSUE = "Please raise an issue on the Python SDK GitHub repository providing context in which this error occurred."
 
 
 class ErrorInfo(Protocol):