test_postgres.py

#  Copyright 2025 Collate
#  Licensed under the Collate Community License, Version 1.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
Test Postgres using the topology
"""

import types
from unittest import TestCase
from unittest.mock import MagicMock, patch

from sqlalchemy.types import VARCHAR

from metadata.generated.schema.entity.data.database import Database
from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
from metadata.generated.schema.entity.data.table import (
    Column,
    Constraint,
    DataType,
    TableType,
)
from metadata.generated.schema.entity.services.databaseService import (
    DatabaseConnection,
    DatabaseService,
    DatabaseServiceType,
)
from metadata.generated.schema.metadataIngestion.workflow import (
    OpenMetadataWorkflowConfig,
)
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.generated.schema.type.filterPattern import FilterPattern
from metadata.ingestion.source.database.common_pg_mappings import (
    GEOMETRY,
    POINT,
    POLYGON,
)
from metadata.ingestion.source.database.postgres.metadata import PostgresSource
from metadata.ingestion.source.database.postgres.usage import PostgresUsageSource
from metadata.ingestion.source.database.postgres.utils import get_postgres_version

mock_postgres_config = {
    "source": {
        "type": "postgres",
        "serviceName": "local_postgres1",
        "serviceConnection": {
            "config": {
                "type": "Postgres",
                "username": "username",
                "authType": {
                    "password": "password",
                },
                "hostPort": "localhost:5432",
                "database": "postgres",
                "sslMode": "verify-ca",
                "sslConfig": {
                    "caCertificate": "CA certificate content",
                },
            }
        },
        "sourceConfig": {
            "config": {
                "type": "DatabaseMetadata",
            }
        },
    },
    "sink": {
        "type": "metadata-rest",
        "config": {},
    },
    "workflowConfig": {
        "openMetadataServerConfig": {
            "hostPort": "http://localhost:8585/api",
            "authProvider": "openmetadata",
            "securityConfig": {"jwtToken": "postgres"},
        }
    },
}

mock_postgres_usage_config = {
    "source": {
        "type": "postgres-usage",
        "serviceName": "local_postgres1",
        "serviceConnection": {
            "config": {
                "type": "Postgres",
                "username": "username",
                "authType": {
                    "password": "password",
                },
                "hostPort": "localhost:5432",
                "database": "postgres",
            }
        },
        "sourceConfig": {
            "config": {
                "type": "DatabaseUsage",
                "queryLogDuration": 1,
            }
        },
    },
    "sink": {
        "type": "metadata-rest",
        "config": {},
    },
    "workflowConfig": {
        "openMetadataServerConfig": {
            "hostPort": "http://localhost:8585/api",
            "authProvider": "openmetadata",
            "securityConfig": {"jwtToken": "postgres"},
        }
    },
}

mock_postgres_usage_config_custom_source = {
    "source": {
        "type": "postgres-usage",
        "serviceName": "local_postgres1",
        "serviceConnection": {
            "config": {
                "type": "Postgres",
                "username": "username",
                "authType": {
                    "password": "password",
                },
                "hostPort": "localhost:5432",
                "database": "postgres",
                "queryStatementSource": "my_schema.custom_pg_stat_statements",
            }
        },
        "sourceConfig": {
            "config": {
                "type": "DatabaseUsage",
                "queryLogDuration": 1,
            }
        },
    },
    "sink": {
        "type": "metadata-rest",
        "config": {},
    },
    "workflowConfig": {
        "openMetadataServerConfig": {
            "hostPort": "http://localhost:8585/api",
            "authProvider": "openmetadata",
            "securityConfig": {"jwtToken": "postgres"},
        }
    },
}

MOCK_DATABASE_SERVICE = DatabaseService(
    id="85811038-099a-11ed-861d-0242ac120002",
    name="postgres_source",
    connection=DatabaseConnection(),
    serviceType=DatabaseServiceType.Postgres,
)

MOCK_DATABASE = Database(
    id="2aaa012e-099a-11ed-861d-0242ac120002",
    name="118146679784",
    fullyQualifiedName="postgres_source.default",
    displayName="118146679784",
    description="",
    service=EntityReference(
        id="85811038-099a-11ed-861d-0242ac120002",
        type="databaseService",
    ),
)

MOCK_DATABASE_SCHEMA = DatabaseSchema(
    id="2aaa012e-099a-11ed-861d-0242ac120056",
    name="default",
    fullyQualifiedName="postgres_source.118146679784.default",
    displayName="default",
    description="",
    database=EntityReference(
        id="2aaa012e-099a-11ed-861d-0242ac120002",
        type="database",
    ),
    service=EntityReference(
        id="2aaa012e-099a-11ed-861d-0242ac120002",
        type="database",
    ),
)


MOCK_COLUMN_VALUE = [
    {
        "name": "username",
        "type": VARCHAR(),
        "nullable": True,
        "default": None,
        "autoincrement": False,
        "system_data_type": "varchar(50)",
        "comment": None,
    },
    {
        "name": "geom_c",
        "type": GEOMETRY(),
        "nullable": True,
        "default": None,
        "autoincrement": False,
        "system_data_type": "geometry",
        "comment": None,
    },
    {
        "name": "point_c",
        "type": POINT(),
        "nullable": True,
        "default": None,
        "autoincrement": False,
        "system_data_type": "point",
        "comment": None,
    },
    {
        "name": "polygon_c",
        "type": POLYGON(),
        "nullable": True,
        "default": None,
        "autoincrement": False,
        "comment": None,
        "system_data_type": "polygon",
    },
]


EXPECTED_COLUMN_VALUE = [
    Column(
        name="username",
        displayName=None,
        dataType=DataType.VARCHAR,
        arrayDataType=None,
        dataLength=1,
        precision=None,
        scale=None,
        dataTypeDisplay="varchar(50)",
        description=None,
        fullyQualifiedName=None,
        tags=None,
        constraint=Constraint.NULL,
        ordinalPosition=None,
        jsonSchema=None,
        children=None,
        customMetrics=None,
        profile=None,
    ),
    Column(
        name="geom_c",
        displayName=None,
        dataType=DataType.GEOMETRY,
        arrayDataType=None,
        dataLength=1,
        precision=None,
        scale=None,
        dataTypeDisplay="geometry",
        description=None,
        fullyQualifiedName=None,
        tags=None,
        constraint=Constraint.NULL,
        ordinalPosition=None,
        jsonSchema=None,
        children=None,
        customMetrics=None,
        profile=None,
    ),
    Column(
        name="point_c",
        displayName=None,
        dataType=DataType.GEOMETRY,
        arrayDataType=None,
        dataLength=1,
        precision=None,
        scale=None,
        dataTypeDisplay="point",
        description=None,
        fullyQualifiedName=None,
        tags=None,
        constraint=Constraint.NULL,
        ordinalPosition=None,
        jsonSchema=None,
        children=None,
        customMetrics=None,
        profile=None,
    ),
    Column(
        name="polygon_c",
        displayName=None,
        dataType=DataType.GEOMETRY,
        arrayDataType=None,
        dataLength=1,
        precision=None,
        scale=None,
        dataTypeDisplay="polygon",
        description=None,
        fullyQualifiedName=None,
        tags=None,
        constraint=Constraint.NULL,
        ordinalPosition=None,
        jsonSchema=None,
        children=None,
        customMetrics=None,
        profile=None,
    ),
]


class PostgresUnitTest(TestCase):
    @patch(
        "metadata.ingestion.source.database.common_db_source.CommonDbSourceService.test_connection"
    )
    def __init__(self, methodName, test_connection) -> None:
        super().__init__(methodName)
        test_connection.return_value = False
        self.config = OpenMetadataWorkflowConfig.model_validate(mock_postgres_config)
        self.postgres_source = PostgresSource.create(
            mock_postgres_config["source"],
            self.config.workflowConfig.openMetadataServerConfig,
        )

        self.postgres_source.context.get().__dict__[
            "database_service"
        ] = MOCK_DATABASE_SERVICE.name.root
        self.postgres_source.context.get().__dict__[
            "database"
        ] = MOCK_DATABASE.name.root
        self.postgres_source.context.get().__dict__[
            "database_schema"
        ] = MOCK_DATABASE_SCHEMA.name.root

        self.usage_config = OpenMetadataWorkflowConfig.model_validate(
            mock_postgres_usage_config
        )
        with patch(
            "metadata.ingestion.source.database.postgres.usage.PostgresUsageSource.test_connection"
        ):
            self.postgres_usage_source = PostgresUsageSource.create(
                mock_postgres_usage_config["source"],
                self.usage_config.workflowConfig.openMetadataServerConfig,
            )

    def test_datatype(self):
        inspector = types.SimpleNamespace()
        inspector.get_columns = (
            lambda table_name, schema_name, table_type, db_name: MOCK_COLUMN_VALUE
        )
        inspector.get_pk_constraint = lambda table_name, schema_name: []
        inspector.get_unique_constraints = lambda table_name, schema_name: []
        inspector.get_foreign_keys = lambda table_name, schema_name: []
        result, _, _ = self.postgres_source.get_columns_and_constraints(
            "public", "user", "postgres", inspector, TableType.Regular
        )
        for i, _ in enumerate(EXPECTED_COLUMN_VALUE):
            self.assertEqual(result[i], EXPECTED_COLUMN_VALUE[i])

    def test_get_stored_procedures(self):
        """
        Test fetching stored procedures with filter
        """
        self.postgres_source.source_config.includeStoredProcedures = True
        self.postgres_source.source_config.storedProcedureFilterPattern = FilterPattern(
            excludes=["sp_exclude"]
        )
        self.postgres_source.context.get().__dict__["database"] = "test_db"
        self.postgres_source.context.get().__dict__["database_schema"] = "test_schema"

        mock_engine = MagicMock()
        self.postgres_source.engine = mock_engine

        # Mock rows
        row1 = MagicMock()
        row1._mapping = {
            "procedure_name": "sp_include",
            "schema_name": "test_schema",
            "definition": "def1",
            "language": "SQL",
            "procedure_type": "PROCEDURE",
        }
        row2 = MagicMock()
        row2._mapping = {
            "procedure_name": "sp_exclude",
            "schema_name": "test_schema",
            "definition": "def2",
            "language": "SQL",
            "procedure_type": "PROCEDURE",
        }

        # PostgreSQL get_stored_procedures calls _get_stored_procedures_internal twice
        # once for procedures and once for functions
        mock_result_proc = MagicMock()
        mock_result_proc.all.return_value = [row1, row2]

        mock_result_func = MagicMock()
        mock_result_func.all.return_value = []

        mock_conn = MagicMock()
        mock_conn.execute.side_effect = [mock_result_proc, mock_result_func]
        mock_engine.connect.return_value.__enter__ = MagicMock(return_value=mock_conn)
        mock_engine.connect.return_value.__exit__ = MagicMock(return_value=False)

        results = list(self.postgres_source.get_stored_procedures())

        self.assertEqual(len(results), 1)
        self.assertEqual(results[0].name, "sp_include")

    def test_get_version_info(self):
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_engine.connect.return_value.__enter__ = MagicMock(return_value=mock_conn)
        mock_engine.connect.return_value.__exit__ = MagicMock(return_value=False)

        mock_conn.execute.return_value.all.return_value = [["110016"]]
        self.assertEqual("110016", get_postgres_version(mock_engine))

        mock_conn.execute.return_value.all.return_value = [["90624"]]
        self.assertEqual("90624", get_postgres_version(mock_engine))

        mock_conn.execute.return_value.all.return_value = [[]]
        self.assertIsNone(get_postgres_version(mock_engine))

    @patch("sqlalchemy.engine.base.Engine")
    @patch(
        "metadata.ingestion.source.database.common_db_source.CommonDbSourceService.connection"
    )
    def test_close_connection(self, engine, connection):
        connection.return_value = True
        self.postgres_source.close()

    def test_query_statement_source_default(self):
        """Test that default query statement source is pg_stat_statements"""
        # Mock the engine with time column check
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_conn.__enter__ = MagicMock(return_value=mock_conn)
        mock_conn.__exit__ = MagicMock(return_value=None)
        mock_conn.execute.return_value = [("total_exec_time",)]
        mock_engine.connect.return_value = mock_conn
        self.postgres_usage_source.engine = mock_engine

        sql_statement = self.postgres_usage_source.get_sql_statement()
        # Verify default pg_stat_statements is used
        self.assertIn("pg_stat_statements s", sql_statement)

    def test_query_statement_source_custom(self):
        """Test that custom query statement source is used when configured"""
        with patch(
            "metadata.ingestion.source.database.postgres.usage.PostgresUsageSource.test_connection"
        ):
            custom_config = OpenMetadataWorkflowConfig.model_validate(
                mock_postgres_usage_config_custom_source
            )
            custom_usage_source = PostgresUsageSource.create(
                mock_postgres_usage_config_custom_source["source"],
                custom_config.workflowConfig.openMetadataServerConfig,
            )

            # Mock the engine with time column check
            mock_engine = MagicMock()
            mock_conn = MagicMock()
            mock_conn.__enter__ = MagicMock(return_value=mock_conn)
            mock_conn.__exit__ = MagicMock(return_value=None)
            mock_conn.execute.return_value = [("total_exec_time",)]
            mock_engine.connect.return_value = mock_conn
            custom_usage_source.engine = mock_engine

            sql_statement = custom_usage_source.get_sql_statement()
            # Verify custom source is used
            self.assertIn("my_schema.custom_pg_stat_statements s", sql_statement)
            # Verify default source is NOT used (check for the default without the custom schema prefix)
            self.assertNotIn("FROM\n  pg_stat_statements s", sql_statement)

    def test_mark_deleted_schemas_enabled(self):
        """Test mark deleted schemas when the config is enabled"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedSchemas = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database information
        self.postgres_source.context.get().__dict__["database"] = "test_db"
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the schema entity source state
        self.postgres_source.schema_entity_source_state = {"test_schema_fqn"}

        # Mock the _get_filtered_schema_names method
        with patch.object(
            self.postgres_source, "_get_filtered_schema_names"
        ) as mock_filtered_schemas:
            mock_filtered_schemas.return_value = [
                "test_schema_fqn",
                "another_schema_fqn",
            ]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_schemas_as_deleted())

                # Verify that delete_entity_from_source was called with correct parameters
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args
                self.assertEqual(call_args[1]["entity_type"], DatabaseSchema)
                self.assertEqual(call_args[1]["mark_deleted_entity"], True)
                self.assertEqual(
                    call_args[1]["params"], {"database": "test_service.test_db"}
                )

                # Verify the entity_source_state contains both processed and filtered schemas
                expected_source_state = {
                    "test_schema_fqn",
                    "test_schema_fqn",
                    "another_schema_fqn",
                }
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )

    def test_mark_deleted_schemas_disabled(self):
        """Test mark deleted schemas when the config is disabled"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedSchemas = False
        self.postgres_source.source_config = mock_source_config

        # Call the method
        result = list(self.postgres_source.mark_schemas_as_deleted())

        # Verify that no deletion operations are performed
        self.assertEqual(result, [])

    def test_mark_deleted_schemas_no_database_context(self):
        """Test mark deleted schemas when no database is in context"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedSchemas = True
        self.postgres_source.source_config = mock_source_config

        # Remove database from context
        self.postgres_source.context.get().__dict__.pop("database", None)

        # Call the method and expect ValueError
        with self.assertRaises(ValueError) as context:
            list(self.postgres_source.mark_schemas_as_deleted())

        self.assertIn("No Database found in the context", str(context.exception))

    def test_mark_deleted_databases_enabled(self):
        """Test mark deleted databases when the config is enabled"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedDatabases = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database service information
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the database entity source state
        self.postgres_source.database_entity_source_state = {"test_db_fqn"}

        # Mock the _get_filtered_database_names method
        with patch.object(
            self.postgres_source, "_get_filtered_database_names"
        ) as mock_filtered_dbs:
            mock_filtered_dbs.return_value = ["test_db", "another_db"]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_databases_as_deleted())

                # Verify that delete_entity_from_source was called with correct parameters
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args
                self.assertEqual(call_args[1]["entity_type"], Database)
                self.assertEqual(call_args[1]["mark_deleted_entity"], True)
                self.assertEqual(call_args[1]["params"], {"service": "test_service"})

                # Verify the entity_source_state contains both processed and filtered databases
                expected_source_state = {
                    "test_db_fqn",
                    "test_service.test_db",
                    "test_service.another_db",
                }
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )

    def test_mark_deleted_databases_disabled(self):
        """Test mark deleted databases when the config is disabled"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedDatabases = False
        self.postgres_source.source_config = mock_source_config

        # Call the method
        result = list(self.postgres_source.mark_databases_as_deleted())

        # Verify that no deletion operations are performed
        self.assertEqual(result, [])

    def test_mark_deleted_schemas_with_schema_filter_pattern(self):
        """Test mark deleted schemas with schema filter pattern applied"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedSchemas = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database information
        self.postgres_source.context.get().__dict__["database"] = "test_db"
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the schema entity source state with some existing schemas
        self.postgres_source.schema_entity_source_state = {
            "test_service.test_db.schema1",
            "test_service.test_db.schema2",
        }

        # Mock the _get_filtered_schema_names method to return filtered schemas
        with patch.object(
            self.postgres_source, "_get_filtered_schema_names"
        ) as mock_filtered_schemas:
            mock_filtered_schemas.return_value = ["test_service.test_db.schema1"]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_schemas_as_deleted())

                # Verify that delete_entity_from_source was called
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args

                # Verify the entity_source_state contains both processed and filtered schemas
                expected_source_state = {
                    "test_service.test_db.schema1",
                    "test_service.test_db.schema2",
                    "test_service.test_db.schema1",
                }
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )

    def test_mark_deleted_databases_with_database_filter_pattern(self):
        """Test mark deleted databases with database filter pattern applied"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedDatabases = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database service information
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the database entity source state with some existing databases
        self.postgres_source.database_entity_source_state = {
            "test_service.db1",
            "test_service.db2",
        }

        # Mock the _get_filtered_database_names method to return filtered databases
        with patch.object(
            self.postgres_source, "_get_filtered_database_names"
        ) as mock_filtered_dbs:
            mock_filtered_dbs.return_value = ["db1"]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_databases_as_deleted())

                # Verify that delete_entity_from_source was called
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args

                # Verify the entity_source_state contains both processed and filtered databases
                expected_source_state = {
                    "test_service.db1",
                    "test_service.db2",
                    "test_service.db1",
                }
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )

    def test_mark_deleted_schemas_empty_source_state(self):
        """Test mark deleted schemas with empty source state"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedSchemas = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database information
        self.postgres_source.context.get().__dict__["database"] = "test_db"
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock empty schema entity source state
        self.postgres_source.schema_entity_source_state = set()

        # Mock the _get_filtered_schema_names method
        with patch.object(
            self.postgres_source, "_get_filtered_schema_names"
        ) as mock_filtered_schemas:
            mock_filtered_schemas.return_value = ["test_service.test_db.schema1"]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_schemas_as_deleted())

                # Verify that delete_entity_from_source was called with only filtered schemas
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args
                expected_source_state = {"test_service.test_db.schema1"}
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )

    def test_mark_deleted_databases_empty_source_state(self):
        """Test mark deleted databases with empty source state"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedDatabases = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database service information
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock empty database entity source state
        self.postgres_source.database_entity_source_state = set()

        # Mock the _get_filtered_database_names method
        with patch.object(
            self.postgres_source, "_get_filtered_database_names"
        ) as mock_filtered_dbs:
            mock_filtered_dbs.return_value = ["db1"]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_databases_as_deleted())

                # Verify that delete_entity_from_source was called with only filtered databases
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args
                expected_source_state = {"test_service.db1"}
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )

    def test_mark_deleted_schemas_exception_handling(self):
        """Test mark deleted schemas exception handling"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedSchemas = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database information
        self.postgres_source.context.get().__dict__["database"] = "test_db"
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the schema entity source state
        self.postgres_source.schema_entity_source_state = {"test_schema_fqn"}

        # Mock the _get_filtered_schema_names method to raise an exception
        with patch.object(
            self.postgres_source, "_get_filtered_schema_names"
        ) as mock_filtered_schemas:
            mock_filtered_schemas.side_effect = Exception("Test exception")

            # Call the method and expect it to handle the exception gracefully
            with self.assertRaises(Exception):
                list(self.postgres_source.mark_schemas_as_deleted())

    def test_mark_deleted_databases_exception_handling(self):
        """Test mark deleted databases exception handling"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedDatabases = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database service information
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the database entity source state
        self.postgres_source.database_entity_source_state = {"test_db_fqn"}

        # Mock the _get_filtered_database_names method to raise an exception
        with patch.object(
            self.postgres_source, "_get_filtered_database_names"
        ) as mock_filtered_dbs:
            mock_filtered_dbs.side_effect = Exception("Test exception")

            # Call the method and expect it to handle the exception gracefully
            with self.assertRaises(Exception):
                list(self.postgres_source.mark_databases_as_deleted())

    def test_mark_deleted_schemas_with_multiple_schemas(self):
        """Test mark deleted schemas with multiple schemas in source state"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedSchemas = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database information
        self.postgres_source.context.get().__dict__["database"] = "test_db"
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the schema entity source state with multiple schemas
        self.postgres_source.schema_entity_source_state = {
            "test_service.test_db.schema1",
            "test_service.test_db.schema2",
            "test_service.test_db.schema3",
        }

        # Mock the _get_filtered_schema_names method
        with patch.object(
            self.postgres_source, "_get_filtered_schema_names"
        ) as mock_filtered_schemas:
            mock_filtered_schemas.return_value = [
                "test_service.test_db.schema1",
                "test_service.test_db.schema2",
            ]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_schemas_as_deleted())

                # Verify that delete_entity_from_source was called
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args

                # Verify the entity_source_state contains all schemas
                expected_source_state = {
                    "test_service.test_db.schema1",
                    "test_service.test_db.schema2",
                    "test_service.test_db.schema3",
                    "test_service.test_db.schema1",
                    "test_service.test_db.schema2",
                }
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )

    def test_mark_deleted_databases_with_multiple_databases(self):
        """Test mark deleted databases with multiple databases in source state"""
        # Create a mock source config with the required attributes
        mock_source_config = MagicMock()
        mock_source_config.markDeletedDatabases = True
        self.postgres_source.source_config = mock_source_config

        # Mock the context to have database service information
        self.postgres_source.context.get().__dict__["database_service"] = "test_service"

        # Mock the database entity source state with multiple databases
        self.postgres_source.database_entity_source_state = {
            "test_service.db1",
            "test_service.db2",
            "test_service.db3",
        }

        # Mock the _get_filtered_database_names method
        with patch.object(
            self.postgres_source, "_get_filtered_database_names"
        ) as mock_filtered_dbs:
            mock_filtered_dbs.return_value = ["db1", "db2"]

            # Mock the delete_entity_from_source function
            with patch(
                "metadata.ingestion.source.database.database_service.delete_entity_from_source"
            ) as mock_delete:
                mock_delete.return_value = iter([])

                # Call the method
                result = list(self.postgres_source.mark_databases_as_deleted())

                # Verify that delete_entity_from_source was called
                mock_delete.assert_called_once()
                call_args = mock_delete.call_args

                # Verify the entity_source_state contains all databases
                expected_source_state = {
                    "test_service.db1",
                    "test_service.db2",
                    "test_service.db3",
                    "test_service.db1",
                    "test_service.db2",
                }
                self.assertEqual(
                    call_args[1]["entity_source_state"], expected_source_state
                )


class TestPostgresCommonMappings(TestCase):
    """Verify extended type entries in the shared PostgreSQL ischema_names map."""

    def test_tid_type_registered(self):
        """'tid' must be present in the PostgreSQL ischema_names after common_pg_mappings is loaded."""
        # common_pg_mappings registers types as a side-effect of module import
        from sqlalchemy.dialects.postgresql.base import (
            ischema_names as pg_ischema_names,
        )

        import metadata.ingestion.source.database.common_pg_mappings  # noqa: F401

        self.assertIn("tid", pg_ischema_names)

    def test_tid_maps_to_string(self):
        """'tid' must map to a String-compatible SQLAlchemy type."""
        from sqlalchemy import String as SqlAlchemyString
        from sqlalchemy.dialects.postgresql.base import (
            ischema_names as pg_ischema_names,
        )

        import metadata.ingestion.source.database.common_pg_mappings  # noqa: F401

        tid_type = pg_ischema_names["tid"]
        self.assertIs(tid_type, SqlAlchemyString)