diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py b/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py index 8104fa2e768..4d329250bb2 100644 --- a/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/postgres/metadata.py @@ -60,6 +60,7 @@ from metadata.ingestion.source.database.postgres.models import PostgresStoredPro from metadata.ingestion.source.database.postgres.queries import ( POSTGRES_GET_ALL_TABLE_PG_POLICY, POSTGRES_GET_DB_NAMES, + POSTGRES_GET_FUNCTIONS, POSTGRES_GET_STORED_PROCEDURES, POSTGRES_GET_TABLE_NAMES, POSTGRES_PARTITION_DETAILS, @@ -273,25 +274,33 @@ class PostgresSource(CommonDbSourceService, MultiDBSource): ) ) + def _get_stored_procedures_internal( + self, query: str + ) -> Iterable[PostgresStoredProcedure]: + results = self.engine.execute(query).all() + for row in results: + try: + stored_procedure = PostgresStoredProcedure.model_validate( + dict(row._mapping) + ) + yield stored_procedure + except Exception as exc: + logger.error() + self.status.failed( + error=StackTraceError( + name=dict(row).get("name", "UNKNOWN"), + error=f"Error parsing Stored Procedure payload: {exc}", + stackTrace=traceback.format_exc(), + ) + ) + def get_stored_procedures(self) -> Iterable[PostgresStoredProcedure]: """List stored procedures""" if self.source_config.includeStoredProcedures: - results = self.engine.execute(POSTGRES_GET_STORED_PROCEDURES).all() - for row in results: - try: - stored_procedure = PostgresStoredProcedure.model_validate( - dict(row._mapping) - ) - yield stored_procedure - except Exception as exc: - logger.error() - self.status.failed( - error=StackTraceError( - name=dict(row).get("name", "UNKNOWN"), - error=f"Error parsing Stored Procedure payload: {exc}", - stackTrace=traceback.format_exc(), - ) - ) + yield from self._get_stored_procedures_internal( + POSTGRES_GET_STORED_PROCEDURES + ) + yield from self._get_stored_procedures_internal(POSTGRES_GET_FUNCTIONS) def yield_stored_procedure( self, stored_procedure @@ -312,6 +321,7 @@ class PostgresSource(CommonDbSourceService, MultiDBSource): database_name=self.context.get().database, schema_name=self.context.get().database_schema, ), + storedProcedureType=stored_procedure.procedure_type, ) yield Either(right=stored_procedure_request) self.register_record_stored_proc_request(stored_procedure_request) diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/models.py b/ingestion/src/metadata/ingestion/source/database/postgres/models.py index a319237b570..e06cb3d9a1f 100644 --- a/ingestion/src/metadata/ingestion/source/database/postgres/models.py +++ b/ingestion/src/metadata/ingestion/source/database/postgres/models.py @@ -23,3 +23,4 @@ class PostgresStoredProcedure(BaseModel): schema: str = Field(alias="schema_name") definition: str language: Optional[str] = None + procedure_type: Optional[str] = Field(None, alias="procedure_type") diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/queries.py b/ingestion/src/metadata/ingestion/source/database/postgres/queries.py index 03a66ede0a3..af21a4d0db1 100644 --- a/ingestion/src/metadata/ingestion/source/database/postgres/queries.py +++ b/ingestion/src/metadata/ingestion/source/database/postgres/queries.py @@ -217,8 +217,25 @@ POSTGRES_GET_STORED_PROCEDURES = """ nspname AS schema_name, proargtypes AS argument_types, prorettype::regtype AS return_type, - prosrc AS definition + prosrc AS definition, + 'StoredProcedure' as procedure_type FROM pg_proc JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid WHERE prokind = 'p'; """ + +POSTGRES_GET_FUNCTIONS = """ +SELECT + proname AS procedure_name, + nspname AS schema_name, + proargtypes AS argument_types, + prorettype :: regtype AS return_type, + prosrc AS definition, + 'Function' as procedure_type +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'); +""" diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/data/storedProcedure.json b/openmetadata-spec/src/main/resources/json/schema/entity/data/storedProcedure.json index cc7c1cb76cf..a46a957ce9b 100644 --- a/openmetadata-spec/src/main/resources/json/schema/entity/data/storedProcedure.json +++ b/openmetadata-spec/src/main/resources/json/schema/entity/data/storedProcedure.json @@ -6,7 +6,9 @@ "description": "A `StoredProcedure` entity that contains the set of code statements with an assigned name and is defined in a `Database Schema`.\"", "type": "object", "javaType": "org.openmetadata.schema.entity.data.StoredProcedure", - "javaInterfaces": ["org.openmetadata.schema.EntityInterface"], + "javaInterfaces": [ + "org.openmetadata.schema.EntityInterface" + ], "definitions": { "storedProcedureType": { "javaType": "org.openmetadata.schema.type.StoredProcedureType", @@ -16,7 +18,8 @@ "enum": [ "StoredProcedure", "UDF", - "StoredPackage" + "StoredPackage", + "Function" ], "javaEnums": [ { @@ -27,6 +30,9 @@ }, { "name": "StoredPackage" + }, + { + "name": "Function" } ] }, diff --git a/openmetadata-ui/src/main/resources/ui/src/generated/entity/data/storedProcedure.ts b/openmetadata-ui/src/main/resources/ui/src/generated/entity/data/storedProcedure.ts index 7ce097e43cf..1f1e63add22 100644 --- a/openmetadata-ui/src/main/resources/ui/src/generated/entity/data/storedProcedure.ts +++ b/openmetadata-ui/src/main/resources/ui/src/generated/entity/data/storedProcedure.ts @@ -501,6 +501,7 @@ export enum Language { * This schema defines the type of the type of Procedures */ export enum StoredProcedureType { + Function = "Function", StoredPackage = "StoredPackage", StoredProcedure = "StoredProcedure", Udf = "UDF",