Optimize ai prompts with the 4-d method

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

@@ -226,6 +226,15 @@ def yield_database(
             owners=self.get_database_owner_ref(database_name),
         )
 
+        # Store database owner in context for schema/table inheritance
+        database_owner_ref = self.get_database_owner_ref(database_name)
+        if database_owner_ref and database_owner_ref.root:
+            database_owner_name = database_owner_ref.root[0].name
+            self.context.get().upsert("database_owner", database_owner_name)
+        else:
+            # Clear context to avoid residual owner from previous database
+            self.context.get().upsert("database_owner", None)
+
         yield Either(right=database_request)
         self.register_record_database_request(database_request=database_request)
 
@@ -282,6 +291,15 @@ def yield_database_schema(
             owners=self.get_schema_owner_ref(schema_name),
         )
 
+        # Store schema owner in context for table inheritance
+        schema_owner_ref = self.get_schema_owner_ref(schema_name)
+        if schema_owner_ref and schema_owner_ref.root:
+            schema_owner_name = schema_owner_ref.root[0].name
+            self.context.get().upsert("schema_owner", schema_owner_name)
+        else:
+            # Clear schema_owner if not present, tables will inherit from database_owner
+            self.context.get().upsert("schema_owner", None)
+
         yield Either(right=schema_request)
         self.register_record_schema_request(schema_request=schema_request)
 

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

@@ -633,12 +633,8 @@ def get_schema_owner_ref(self, schema_name: str) -> Optional[EntityReferenceList
             EntityReferenceList with owner or None
         """
         try:
-            parent_owner = None
-            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
+            # Read database_owner directly from context
+            parent_owner = getattr(self.context.get(), "database_owner", None)
 
             schema_fqn = f"{self.context.get().database}.{schema_name}"
 
@@ -678,14 +674,10 @@ def get_owner_ref(self, table_name: str) -> Optional[EntityReferenceList]:
             EntityReferenceList with owner or None
         """
         try:
-            parent_owner = None
-            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
+            # Prioritize schema_owner, fallback to database_owner for inheritance
+            parent_owner = getattr(self.context.get(), "schema_owner", None)
+            if not parent_owner:
+                parent_owner = getattr(self.context.get(), "database_owner", None)
 
             table_fqn = f"{self.context.get().database}.{self.context.get().database_schema}.{table_name}"