From b4cd7b7046bd630a8b5f117e34668e3a967e7040 Mon Sep 17 00:00:00 2001
From: Suman Maharana <sumanmaharana786@gmail.com>
Date: Mon, 30 Jun 2025 10:39:09 +0530
Subject: [PATCH] Add: Postgres SP and UDF descriptions (#22021)

---
 .../source/database/postgres/metadata.py      | 20 +++++++++++++++----
 .../source/database/postgres/models.py        |  1 +
 .../source/database/postgres/queries.py       | 11 ++++++----
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py b/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py
index 1479d24e4f1..69b24786d4e 100644
--- a/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py
+++ b/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py
@@ -40,7 +40,11 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    Markdown,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -300,9 +304,15 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
         """List stored procedures"""
         if self.source_config.includeStoredProcedures:
             yield from self._get_stored_procedures_internal(
-                POSTGRES_GET_STORED_PROCEDURES
+                POSTGRES_GET_STORED_PROCEDURES.format(
+                    schema_name=self.context.get().database_schema
+                )
+            )
+            yield from self._get_stored_procedures_internal(
+                POSTGRES_GET_FUNCTIONS.format(
+                    schema_name=self.context.get().database_schema
+                )
             )
-            yield from self._get_stored_procedures_internal(POSTGRES_GET_FUNCTIONS)
 
     def yield_stored_procedure(
         self, stored_procedure
@@ -311,7 +321,9 @@ class PostgresSource(CommonDbSourceService, MultiDBSource):
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
                 name=EntityName(stored_procedure.name),
-                description=None,
+                description=Markdown(stored_procedure.description)
+                if stored_procedure.description
+                else None,
                 storedProcedureCode=StoredProcedureCode(
                     language=STORED_PROC_LANGUAGE_MAP.get(stored_procedure.language),
                     code=stored_procedure.definition,
diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/models.py b/ingestion/src/metadata/ingestion/source/database/postgres/models.py
index a8be7fe6b4e..83efa9c7306 100644
--- a/ingestion/src/metadata/ingestion/source/database/postgres/models.py
+++ b/ingestion/src/metadata/ingestion/source/database/postgres/models.py
@@ -24,3 +24,4 @@ class PostgresStoredProcedure(BaseModel):
     definition: str
     language: Optional[str] = None
     procedure_type: Optional[str] = Field(None, alias="procedure_type")
+    description: Optional[str] = Field(None, alias="description")
diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/queries.py b/ingestion/src/metadata/ingestion/source/database/postgres/queries.py
index eb68d1a6767..23c0460eb93 100644
--- a/ingestion/src/metadata/ingestion/source/database/postgres/queries.py
+++ b/ingestion/src/metadata/ingestion/source/database/postgres/queries.py
@@ -225,10 +225,12 @@ POSTGRES_GET_STORED_PROCEDURES = """
         proargtypes AS argument_types,
         prorettype::regtype AS return_type,
         prosrc AS definition,
-        'StoredProcedure' as procedure_type
+        'StoredProcedure' as procedure_type,
+        obj_description(pg_proc.oid, 'pg_proc') AS description
     FROM pg_proc
     JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
-    WHERE prokind = 'p';
+    WHERE prokind = 'p'
+    and pg_namespace.nspname = '{schema_name}';
 """
 
 POSTGRES_GET_FUNCTIONS = """
@@ -238,11 +240,12 @@ SELECT
     proargtypes AS argument_types,
     prorettype :: regtype AS return_type,
     prosrc AS definition,
-    'Function' as procedure_type
+    'Function' as procedure_type,
+    obj_description(pg_proc.oid, 'pg_proc') AS description
 FROM
     pg_proc
     JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid
 WHERE
     prokind = 'f'
-    and pg_namespace.nspname NOT IN ('pg_catalog', 'information_schema');
+    and pg_namespace.nspname = '{schema_name}';
 """