MINOR: Fix public schema lieage for postgres (#20548)

This commit is contained in:
Mayur Singal 2025-04-02 15:30:24 +05:30 committed by GitHub
parent 7402feba6f
commit c16b3df547
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 2 deletions

View File

@ -31,6 +31,7 @@ from metadata.ingestion.source.database.postgres.queries import POSTGRES_SQL_STA
from metadata.ingestion.source.database.postgres.query_parser import ( from metadata.ingestion.source.database.postgres.query_parser import (
PostgresQueryParserSource, PostgresQueryParserSource,
) )
from metadata.utils.db_utils import PUBLIC_SCHEMA
from metadata.utils.logger import ingestion_logger from metadata.utils.logger import ingestion_logger
logger = ingestion_logger() logger = ingestion_logger()
@ -89,7 +90,7 @@ class PostgresLineageSource(PostgresQueryParserSource, LineageSource):
aborted=self.get_aborted_status(row), aborted=self.get_aborted_status(row),
databaseName=self.get_database_name(row), databaseName=self.get_database_name(row),
serviceName=self.config.serviceName, serviceName=self.config.serviceName,
databaseSchema=self.get_schema_name(row), databaseSchema=self.get_schema_name(row) or PUBLIC_SCHEMA,
duration=row.get("duration"), duration=row.get("duration"),
) )
except Exception as err: except Exception as err:

View File

@ -18,6 +18,9 @@ from typing import Iterable
from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
from metadata.generated.schema.entity.data.table import Table from metadata.generated.schema.entity.data.table import Table
from metadata.generated.schema.entity.services.databaseService import (
DatabaseServiceType,
)
from metadata.generated.schema.type.entityLineage import Source as LineageSource from metadata.generated.schema.type.entityLineage import Source as LineageSource
from metadata.ingestion.api.models import Either from metadata.ingestion.api.models import Either
from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
@ -33,6 +36,8 @@ from metadata.utils.logger import utils_logger
logger = utils_logger() logger = utils_logger()
PUBLIC_SCHEMA = "public"
def get_host_from_host_port(uri: str) -> str: def get_host_from_host_port(uri: str) -> str:
""" """
@ -64,7 +69,7 @@ def get_view_lineage(
schema_name=schema_name, schema_name=schema_name,
table_name=table_name, table_name=table_name,
) )
table_entity = metadata.get_by_name( table_entity: Table = metadata.get_by_name(
entity=Table, entity=Table,
fqn=table_fqn, fqn=table_fqn,
) )
@ -79,6 +84,11 @@ def get_view_lineage(
lineage_parser = LineageParser( lineage_parser = LineageParser(
view_definition, dialect, timeout_seconds=timeout_seconds view_definition, dialect, timeout_seconds=timeout_seconds
) )
if table_entity.serviceType == DatabaseServiceType.Postgres:
# For Postgres, if schema is not defined, we need to use the public schema
schema_name = PUBLIC_SCHEMA
if lineage_parser.source_tables and lineage_parser.target_tables: if lineage_parser.source_tables and lineage_parser.target_tables:
yield from get_lineage_by_query( yield from get_lineage_by_query(
metadata, metadata,