Refactor: Use context for owner inheritance, remove cache

Co-authored-by: yourton.ma <yourton.ma@gmail.com>

Author: cursoragent

ingestion/src/metadata/ingestion/api/topology_runner.py

@@ -446,6 +446,20 @@ def yield_and_update_context(
                 )
 
         self.context.get().update_context_name(stage=stage, right=right)
+        
+        # Store the entity in context for downstream inheritance
+        # After yielding the request to sink, fetch the entity from API to get the complete entity
+        # with all resolved fields (like owners) for child entities to inherit from
+        if stage.context and not same_fingerprint:
+            # Fetch the entity after it has been created/updated by the sink
+            entity = self.metadata.get_by_name(
+                entity=stage.type_,
+                fqn=entity_fqn,
+                fields=["owners"],  # Only fetch owners field for performance
+            )
+            if entity:
+                entity_context_key = f"{stage.context}_entity"
+                self.context.get().upsert(key=entity_context_key, value=entity)
 
     @yield_and_update_context.register
     def _(

ingestion/src/metadata/ingestion/source/database/database_service.py

@@ -225,8 +225,6 @@ class DatabaseServiceSource(
     stored_procedure_source_state: Set = set()
     database_entity_source_state: Set = set()
     schema_entity_source_state: Set = set()
-    # Cache for schema owners to avoid repeated API calls during table ingestion
-    schema_owner_cache: Dict[str, Optional[str]] = {}
     # Big union of types we want to fetch dynamically
     service_connection: DatabaseConnection.model_fields["config"].annotation
 
@@ -598,7 +596,6 @@ def get_database_owner_ref(
             EntityReferenceList with owner or None
         """
         try:
-            # Priority 1: Use ownerConfig if configured
             if (
                 hasattr(self.source_config, "ownerConfig")
                 and self.source_config.ownerConfig
@@ -611,9 +608,6 @@ def get_database_owner_ref(
                     parent_owner=None,  # Database is top level
                 )
                 if owner_ref:
-                    # Cache the resolved owner for schema inheritance
-                    if owner_ref.root:
-                        self.schema_owner_cache[database_name] = owner_ref.root[0].name
                     return owner_ref
 
         except Exception as exc:
@@ -638,14 +632,12 @@ def get_schema_owner_ref(self, schema_name: str) -> Optional[EntityReferenceList
             EntityReferenceList with owner or None
         """
         try:
-            # Get parent owner from cached database owner
             parent_owner = None
-            database_name = self.context.get().database
-            if database_name:
-                parent_owner = self.schema_owner_cache.get(database_name)
-                logger.debug(
-                    f"Schema '{schema_name}': Using cached database owner from '{database_name}': {parent_owner}"
-                )
+            database_entity = getattr(self.context.get(), "database_entity", None)
+            if database_entity:
+                db_owners = database_entity.owners
+                if db_owners and db_owners.root:
+                    parent_owner = db_owners.root[0].name
 
             schema_fqn = f"{self.context.get().database}.{schema_name}"
 
@@ -661,9 +653,6 @@ def get_schema_owner_ref(self, schema_name: str) -> Optional[EntityReferenceList
                     parent_owner=parent_owner,
                 )
                 if owner_ref:
-                    # Cache the resolved owner for table inheritance
-                    if owner_ref.root:
-                        self.schema_owner_cache[schema_fqn] = owner_ref.root[0].name
                     return owner_ref
 
         except Exception as exc:
@@ -689,21 +678,13 @@ def get_owner_ref(self, table_name: str) -> Optional[EntityReferenceList]:
         """
         try:
             parent_owner = None
-            
-            # Get parent owner from cached schema owner
-            # The schema owner was cached when the schema was processed
-            if (
-                hasattr(self.source_config, "ownerConfig")
-                and self.source_config.ownerConfig
-            ):
-                schema_fqn = f"{self.context.get().database}.{self.context.get().database_schema}"
-                
-                # Use cached schema owner if available
-                parent_owner = self.schema_owner_cache.get(schema_fqn)
-                
-                logger.debug(
-                    f"Table '{table_name}': Using cached schema owner from '{schema_fqn}': {parent_owner}"
-                )
+            database_schema_entity = getattr(
+                self.context.get(), "database_schema_entity", None
+            )
+            if database_schema_entity:
+                schema_owners = database_schema_entity.owners
+                if schema_owners and schema_owners.root:
+                    parent_owner = schema_owners.root[0].name
 
             table_fqn = f"{self.context.get().database}.{self.context.get().database_schema}.{table_name}"