[generator] Fixed circular imports due to Process related to SQL assets

Author: Aryamanz29

pyatlan/generator/class_generator.py

@@ -223,14 +223,15 @@ class AssetInfo:
     entity_defs_by_name: Dict[str, EntityDef] = {}
     sub_type_names_to_ignore: Set[str] = set()
     is_core_asset: bool = False
-    _ASSETS_REQUIRE_CLIENT = {
+    _ASSETS_REQUIRE_TYPE_CHECKING = {
         "Badge",
         "Collection",
         "Connection",
         "DataProduct",
         "Referenceable",
         "Purpose",
         "DataQualityRule",
+        "Process",
     }
     _ASSETS_REQUIRE_ASYNC_CLIENT = {
         "Badge",
@@ -253,7 +254,20 @@ class AssetInfo:
         "DocumentDBCollection",
         "FlowDataset",
         "DatabricksAIModelVersion",
+        "Fabric",
+        "FabricDashboard",
+        "FabricDataflow",
+        "FabricActivity",
+        "FabricPage",
+        "FabricWorkspace",
+        "FabricDataPipeline",
+        "FabricSemanticModelTable",
+        "FabricSemanticModelTableColumn",
+        "FabricDataflowEntityColumn",
+        "FabricReport",
+        "FabricSemanticModel",
     }
+    _IGNORE_ASSETS = {}
 
     def __init__(self, name: str, entity_def: EntityDef):
         self._name = name
@@ -304,6 +318,19 @@ def imports_for_referenced_assets(self):
         imports = []
 
         for required_asset in self.required_asset_infos:
+            # To avoid circular import issues with DataQualityRule and Column
+            # Though we import Column in data_quality_rule.py for type hinting purposes,
+            if self.name == "DataQualityRule" and required_asset.name == "Column":
+                continue
+            # FIXME: Temporary fix to avoid circular import issues
+            # Process cant have relationship to Procedure, 
+            # BigqueryRoutine, FabricActivity since its a super type is SQL
+            if self.name == "Process" and required_asset.name in (
+                "Procedure",
+                "BigqueryRoutine",
+                "FabricActivity",
+            ):
+                continue
             if not self.is_core_asset and required_asset.is_core_asset:
                 import_statement = f"from .core.{required_asset.module_name} import {required_asset.name} # noqa: E402, F401"
             else:
@@ -357,6 +384,8 @@ def update_required_asset_names(self) -> None:
         relationship_attribute_defs = self.entity_def.relationship_attribute_defs or []
         for attribute in attribute_defs + relationship_attribute_defs:
             type_name = attribute["typeName"].replace("array<", "").replace(">", "")
+            # if type_name in AssetInfo._IGNORE_ASSETS:
+            #     continue
             if type_name == self._name:
                 continue
             if type_name in AssetInfo.super_type_names_to_ignore:
@@ -438,6 +467,9 @@ def create_modules(cls):
             cls.hierarchy_graph, REFERENCEABLE
         ):
             for asset_name in [parent_name] + successors:
+                # if asset_name in cls._IGNORE_ASSETS:
+                #     print(f"Ignoring asset {asset_name}")
+                #     continue
                 asset_info = cls.asset_info_by_name[asset_name]
                 asset_info.order = order
                 order += 1

pyatlan/generator/templates/imports.jinja2

@@ -118,6 +118,12 @@ from pyatlan.model.structs import (
     DatabricksAIModelVersionMetric,
     AssetHistogram,
     AppWorkflowRunStep,
+    SQLProcedureReturn,
+    SQLProcedureArgument,
+    AssetExternalDQMetadata,
+    AssetGCPDataplexMetadata,
+    AssetGCPDataplexAspectMetadata,
+    AssetExternalDQTestRunHistory,
 )
 from pyatlan.utils import (
     init_guid,

pyatlan/generator/templates/module.jinja2

@@ -6,7 +6,7 @@
 
 {% include 'imports.jinja2' %}
 
-{% if asset_info.name not in asset_info._ASSETS_REQUIRE_CLIENT %}
+{% if asset_info.name not in asset_info._ASSETS_REQUIRE_TYPE_CHECKING %}
 from pyatlan.client.atlan import AtlanClient
 {% endif %}
 
@@ -37,12 +37,17 @@ from .process import Process
 
 {{ asset_info.import_super_class }}
 
-{% if asset_info.name in asset_info._ASSETS_REQUIRE_CLIENT %}
+{% if asset_info.name in asset_info._ASSETS_REQUIRE_TYPE_CHECKING %}
 if TYPE_CHECKING:
     from pyatlan.client.atlan import AtlanClient
     {% if asset_info.name == 'DataQualityRule' %}
     from pyatlan.model.assets import Column
     {% endif %}
+    {% if asset_info.name == 'Process' %}
+    from pyatlan.model.assets import Procedure
+    from pyatlan.model.assets import BigqueryRoutine
+    from pyatlan.model.assets import FabricActivity
+    {% endif %}
     {% if asset_info.name in asset_info._ASSETS_REQUIRE_ASYNC_CLIENT %}
     from pyatlan.client.aio.client import AsyncAtlanClient
     from pyatlan.model.aio import AsyncCustomMetadataProxy