openmetadata-ingestion 1.6.7.0__py3-none-any.whl → 1.7.0.0rc1__py3-none-any.whl

metadata/ingestion/source/database/mssql/metadata.py

@@ -30,7 +30,7 @@ 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
+from metadata.generated.schema.type.basic import EntityName, Markdown
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -41,6 +41,9 @@ from metadata.ingestion.source.database.mssql.models import (
 )
 from metadata.ingestion.source.database.mssql.queries import (
     MSSQL_GET_DATABASE,
+    MSSQL_GET_DATABASE_COMMENTS,
+    MSSQL_GET_SCHEMA_COMMENTS,
+    MSSQL_GET_STORED_PROCEDURE_COMMENTS,
     MSSQL_GET_STORED_PROCEDURES,
 )
 from metadata.ingestion.source.database.mssql.utils import (
@@ -94,6 +97,16 @@ class MssqlSource(CommonDbSourceService, MultiDBSource):
     Database metadata from MSSQL Source
     """
 
+    def __init__(
+        self,
+        config,
+        metadata,
+    ):
+        super().__init__(config, metadata)
+        self.schema_desc_map = {}
+        self.database_desc_map = {}
+        self.stored_procedure_desc_map = {}
+
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
@@ -112,12 +125,60 @@ class MssqlSource(CommonDbSourceService, MultiDBSource):
             return self.service_connection.database
         return None
 
+    def set_schema_description_map(self) -> None:
+        self.schema_desc_map.clear()
+        results = self.engine.execute(MSSQL_GET_SCHEMA_COMMENTS).all()
+        self.schema_desc_map = {
+            (row.DATABASE_NAME, row.SCHEMA_NAME): row.COMMENT for row in results
+        }
+
+    def set_database_description_map(self) -> None:
+        self.database_desc_map.clear()
+        results = self.engine.execute(MSSQL_GET_DATABASE_COMMENTS).all()
+        self.database_desc_map = {row.DATABASE_NAME: row.COMMENT for row in results}
+
+    def set_stored_procedure_description_map(self) -> None:
+        self.stored_procedure_desc_map.clear()
+        results = self.engine.execute(MSSQL_GET_STORED_PROCEDURE_COMMENTS).all()
+        self.stored_procedure_desc_map = {
+            (row.DATABASE_NAME, row.SCHEMA_NAME, row.STORED_PROCEDURE): row.COMMENT
+            for row in results
+        }
+
+    def get_schema_description(self, schema_name: str) -> Optional[str]:
+        """
+        Method to fetch the schema description
+        """
+        return self.schema_desc_map.get((self.context.get().database, schema_name))
+
+    def get_database_description(self, database_name: str) -> Optional[str]:
+        """
+        Method to fetch the database description
+        """
+        return self.database_desc_map.get(database_name)
+
+    def get_stored_procedure_description(self, stored_procedure: str) -> Optional[str]:
+        """
+        Method to fetch the stored procedure description
+        """
+        description = self.stored_procedure_desc_map.get(
+            (
+                self.context.get().database,
+                self.context.get().database_schema,
+                stored_procedure,
+            )
+        )
+        return Markdown(description) if description else None
+
     def get_database_names_raw(self) -> Iterable[str]:
         yield from self._execute_database_query(MSSQL_GET_DATABASE)
 
     def get_database_names(self) -> Iterable[str]:
         if not self.config.serviceConnection.root.config.ingestAllDatabases:
             configured_db = self.config.serviceConnection.root.config.database
+            self.set_schema_description_map()
+            self.set_database_description_map()
+            self.set_stored_procedure_description_map()
             self.set_inspector(database_name=configured_db)
             yield configured_db
         else:
@@ -139,6 +200,9 @@ class MssqlSource(CommonDbSourceService, MultiDBSource):
                     continue
 
                 try:
+                    self.set_schema_description_map()
+                    self.set_database_description_map()
+                    self.set_stored_procedure_description_map()
                     self.set_inspector(database_name=new_database)
                     yield new_database
                 except Exception as exc:
@@ -178,7 +242,9 @@ class MssqlSource(CommonDbSourceService, MultiDBSource):
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
                 name=EntityName(stored_procedure.name),
-                description=None,
+                description=self.get_stored_procedure_description(
+                    stored_procedure.name
+                ),
                 storedProcedureCode=StoredProcedureCode(
                     language=STORED_PROC_LANGUAGE_MAP.get(stored_procedure.language),
                     code=stored_procedure.definition,

metadata/ingestion/source/database/mssql/queries.py

@@ -56,6 +56,49 @@ WHERE
 """
 )
 
+MSSQL_GET_DATABASE_COMMENTS = textwrap.dedent(
+    """
+    SELECT 
+    DB_NAME() AS DATABASE_NAME,
+    CAST(ep.value AS NVARCHAR(MAX)) AS COMMENT
+FROM sys.extended_properties ep
+WHERE ep.class = 0  
+AND ep.name = 'MS_Description'
+"""
+)
+
+MSSQL_GET_SCHEMA_COMMENTS = textwrap.dedent(
+    """
+     SELECT 
+    DB_NAME() AS DATABASE_NAME, 
+    s.name AS SCHEMA_NAME, 
+    CAST(ep.value AS NVARCHAR(MAX)) AS COMMENT 
+FROM sys.schemas s
+LEFT JOIN sys.extended_properties ep 
+    ON ep.major_id = s.schema_id
+    AND ep.minor_id = 0 
+    AND ep.class = 3
+    AND ep.name = 'MS_Description'
+    """
+)
+
+MSSQL_GET_STORED_PROCEDURE_COMMENTS = textwrap.dedent(
+    """
+SELECT 
+    DB_NAME() AS DATABASE_NAME,
+    s.name AS SCHEMA_NAME,
+    p.name AS STORED_PROCEDURE,
+    CAST(ep.value AS NVARCHAR(MAX)) AS COMMENT
+FROM sys.procedures p
+JOIN sys.schemas s ON p.schema_id = s.schema_id
+LEFT JOIN sys.extended_properties ep 
+    ON ep.major_id = p.object_id 
+    AND ep.minor_id = 0 
+    AND ep.class = 1
+    AND ep.name = 'MS_Description';
+"""
+)
+
 MSSQL_ALL_VIEW_DEFINITIONS = textwrap.dedent(
     """
 SELECT

metadata/ingestion/source/database/mysql/utils.py

@@ -12,6 +12,8 @@
 """
 MySQL SQLAlchemy Helper Methods
 """
+
+
 # pylint: disable=protected-access,too-many-branches,too-many-statements,too-many-locals
 from sqlalchemy import util
 from sqlalchemy.dialects.mysql.enumerated import ENUM, SET
@@ -136,9 +138,9 @@ def parse_column(self, line, state):
     raw_type = get_display_datatype(
         col_type=type_,
         char_len=type_instance.length if hasattr(type_instance, "length") else None,
-        precision=type_instance.precision
-        if hasattr(type_instance, "precision")
-        else None,
+        precision=(
+            type_instance.precision if hasattr(type_instance, "precision") else None
+        ),
         scale=type_instance.scale if hasattr(type_instance, "scale") else None,
     )
 

metadata/ingestion/source/database/oracle/connection.py

@@ -41,7 +41,11 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.oracle.queries import CHECK_ACCESS_TO_ALL
+from metadata.ingestion.source.database.oracle.queries import (
+    CHECK_ACCESS_TO_ALL,
+    ORACLE_GET_SCHEMA,
+    ORACLE_GET_STORED_PACKAGES,
+)
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
 
@@ -131,6 +135,12 @@ def get_connection(connection: OracleConnection) -> Engine:
     )
 
 
+class OraclePackageAccessError(Exception):
+    """
+    Raised when unable to access Oracle stored packages
+    """
+
+
 def test_connection(
     metadata: OpenMetadata,
     engine: Engine,
@@ -143,7 +153,19 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    test_conn_queries = {"CheckAccess": CHECK_ACCESS_TO_ALL}
+    def test_oracle_package_access(engine):
+        try:
+            schema_name = engine.execute(ORACLE_GET_SCHEMA).scalar()
+            return ORACLE_GET_STORED_PACKAGES.format(schema=schema_name)
+        except Exception as e:
+            raise OraclePackageAccessError(
+                f"Failed to access Oracle stored packages: {e}"
+            )
+
+    test_conn_queries = {
+        "CheckAccess": CHECK_ACCESS_TO_ALL,
+        "PackageAccess": test_oracle_package_access(engine),
+    }
 
     return test_connection_db_common(
         metadata=metadata,

metadata/ingestion/source/database/oracle/metadata.py

@@ -24,6 +24,7 @@ from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.storedProcedure import (
     Language,
     StoredProcedureCode,
+    StoredProcedureType,
 )
 from metadata.generated.schema.entity.data.table import TableType
 from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
@@ -45,10 +46,11 @@ from metadata.ingestion.source.database.common_db_source import (
     TableNameAndType,
 )
 from metadata.ingestion.source.database.oracle.models import (
-    FetchProcedureList,
-    OracleStoredProcedure,
+    FetchObjectList,
+    OracleStoredObject,
 )
 from metadata.ingestion.source.database.oracle.queries import (
+    ORACLE_GET_STORED_PACKAGES,
     ORACLE_GET_STORED_PROCEDURES,
 )
 from metadata.ingestion.source.database.oracle.utils import (
@@ -181,41 +183,51 @@ class OracleSource(CommonDbSourceService):
             logger.warning(f"Failed to fetch Schema definition for {table_name}: {exc}")
         return None
 
-    def process_result(self, data: FetchProcedureList):
+    def process_result(self, data: FetchObjectList):
         """Process data as per our stored procedure format"""
         result_dict = {}
 
         for row in data:
-            owner, name, line, text = row
+
+            owner, name, line, text, procedure_type = row
             key = (owner, name)
             if key not in result_dict:
-                result_dict[key] = {"lines": [], "text": ""}
+                result_dict[key] = {"lines": [], "text": "", "procedure_type": ""}
             result_dict[key]["lines"].append(line)
             result_dict[key]["text"] += text
+            result_dict[key]["procedure_type"] = procedure_type
 
         # Return the concatenated text for each procedure name, ordered by line
         return result_dict
 
-    def get_stored_procedures(self) -> Iterable[OracleStoredProcedure]:
+    def _get_stored_procedures_internal(
+        self, query: str
+    ) -> Iterable[OracleStoredObject]:
+        results: FetchObjectList = self.engine.execute(
+            query.format(schema=self.context.get().database_schema.upper())
+        ).all()
+        results = self.process_result(data=results)
+        for row in results.items():
+            stored_procedure = OracleStoredObject(
+                name=row[0][1],
+                definition=row[1]["text"],
+                owner=row[0][0],
+                procedure_type=row[1]["procedure_type"],
+            )
+            yield stored_procedure
+
+    def get_stored_procedures(self) -> Iterable[OracleStoredObject]:
         """List Oracle Stored Procedures"""
         if self.source_config.includeStoredProcedures:
-            results: FetchProcedureList = self.engine.execute(
-                ORACLE_GET_STORED_PROCEDURES.format(
-                    schema=self.context.get().database_schema.upper()
-                )
-            ).all()
-            results = self.process_result(data=results)
-            for row in results.items():
-                stored_procedure = OracleStoredProcedure(
-                    name=row[0][1], definition=row[1]["text"], owner=row[0][0]
-                )
-                yield stored_procedure
+            yield from self._get_stored_procedures_internal(
+                ORACLE_GET_STORED_PROCEDURES
+            )
+            yield from self._get_stored_procedures_internal(ORACLE_GET_STORED_PACKAGES)
 
     def yield_stored_procedure(
-        self, stored_procedure: OracleStoredProcedure
+        self, stored_procedure: OracleStoredObject
     ) -> Iterable[Either[CreateStoredProcedureRequest]]:
         """Prepare the stored procedure payload"""
-
         try:
             stored_procedure_request = CreateStoredProcedureRequest(
                 name=EntityName(stored_procedure.name),
@@ -223,6 +235,11 @@ class OracleSource(CommonDbSourceService):
                     language=Language.SQL,
                     code=stored_procedure.definition,
                 ),
+                storedProcedureType=(
+                    StoredProcedureType.StoredPackage
+                    if stored_procedure.procedure_type == "StoredPackage"
+                    else StoredProcedureType.StoredProcedure
+                ),
                 owners=self.metadata.get_reference_by_name(
                     name=stored_procedure.owner.lower(), is_owner=True
                 ),

metadata/ingestion/source/database/oracle/models.py

@@ -6,7 +6,7 @@ from typing import List, Optional
 from pydantic import BaseModel, Field
 
 
-class OracleStoredProcedure(BaseModel):
+class OracleStoredObject(BaseModel):
     """Oracle Stored Procedure list query results"""
 
     name: str
@@ -15,9 +15,10 @@ class OracleStoredProcedure(BaseModel):
         None, description="Will only be informed for non-SQL routines."
     )
     owner: str
+    procedure_type: Optional[str] = Field(None, alias="procedure_type")
 
 
-class FetchProcedure(BaseModel):
+class FetchObject(BaseModel):
     """Oracle Fetch Stored Procedure Raw Model"""
 
     owner: Optional[str] = None
@@ -26,5 +27,5 @@ class FetchProcedure(BaseModel):
     text: str
 
 
-class FetchProcedureList(BaseModel):
-    __name__: List[FetchProcedure]
+class FetchObjectList(BaseModel):
+    __name__: List[FetchObject]

metadata/ingestion/source/database/oracle/queries.py

@@ -87,13 +87,34 @@ SELECT
     OWNER,
     NAME,
     LINE,
-    TEXT
+    TEXT,
+    'StoredProcedure' as procedure_type
 FROM
     DBA_SOURCE
 WHERE
     type = 'PROCEDURE' and owner = '{schema}'
 """
 )
+ORACLE_GET_SCHEMA = """
+    SELECT USERNAME AS SCHEMA_NAME 
+    FROM ALL_USERS 
+    WHERE ROWNUM = 1 
+    ORDER BY USERNAME
+"""
+ORACLE_GET_STORED_PACKAGES = textwrap.dedent(
+    """
+SELECT
+    OWNER,
+    NAME,
+    LINE,
+    TEXT,
+    'StoredPackage' as procedure_type
+
+FROM
+    DBA_SOURCE
+WHERE TYPE IN ('PACKAGE', 'PACKAGE BODY') AND owner = '{schema}'
+"""
+)
 CHECK_ACCESS_TO_ALL = "SELECT table_name FROM DBA_TABLES where ROWNUM < 2"
 ORACLE_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent(
     """

metadata/ingestion/source/database/pinotdb/service_spec.py

@@ -1,4 +1,8 @@
+from metadata.ingestion.source.database.pinotdb.lineage import PinotdbLineageSource
 from metadata.ingestion.source.database.pinotdb.metadata import PinotdbSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=PinotdbSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=PinotdbSource,
+    lineage_source_class=PinotdbLineageSource,
+)

metadata/ingestion/source/database/postgres/metadata.py

@@ -15,9 +15,8 @@ import traceback
 from collections import namedtuple
 from typing import Iterable, Optional, Tuple
 
-from sqlalchemy import String as SqlAlchemyString
 from sqlalchemy import sql
-from sqlalchemy.dialects.postgresql.base import PGDialect, ischema_names
+from sqlalchemy.dialects.postgresql.base import PGDialect
 from sqlalchemy.engine import Inspector
 
 from metadata.generated.schema.api.data.createStoredProcedure import (
@@ -46,17 +45,22 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.column_type_parser import create_sqlalchemy_type
 from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.common_pg_mappings import (
+    INTERVAL_TYPE_MAP,
+    RELKIND_MAP,
+    ischema_names,
+)
 from metadata.ingestion.source.database.mssql.models import STORED_PROC_LANGUAGE_MAP
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.postgres.models import PostgresStoredProcedure
 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,
@@ -95,43 +99,6 @@ TableKey = namedtuple("TableKey", ["schema", "table_name"])
 logger = ingestion_logger()
 
 
-INTERVAL_TYPE_MAP = {
-    "list": PartitionIntervalTypes.COLUMN_VALUE,
-    "hash": PartitionIntervalTypes.COLUMN_VALUE,
-    "range": PartitionIntervalTypes.TIME_UNIT,
-}
-
-RELKIND_MAP = {
-    "r": TableType.Regular,
-    "p": TableType.Partitioned,
-    "f": TableType.Foreign,
-}
-
-GEOMETRY = create_sqlalchemy_type("GEOMETRY")
-POINT = create_sqlalchemy_type("POINT")
-POLYGON = create_sqlalchemy_type("POLYGON")
-
-ischema_names.update(
-    {
-        "geometry": GEOMETRY,
-        "point": POINT,
-        "polygon": POLYGON,
-        "box": create_sqlalchemy_type("BOX"),
-        "bpchar": SqlAlchemyString,
-        "circle": create_sqlalchemy_type("CIRCLE"),
-        "line": create_sqlalchemy_type("LINE"),
-        "lseg": create_sqlalchemy_type("LSEG"),
-        "path": create_sqlalchemy_type("PATH"),
-        "pg_lsn": create_sqlalchemy_type("PG_LSN"),
-        "pg_snapshot": create_sqlalchemy_type("PG_SNAPSHOT"),
-        "tsquery": create_sqlalchemy_type("TSQUERY"),
-        "txid_snapshot": create_sqlalchemy_type("TXID_SNAPSHOT"),
-        "xid": SqlAlchemyString,
-        "xml": create_sqlalchemy_type("XML"),
-    }
-)
-
-
 PGDialect.get_all_table_comments = get_all_table_comments
 PGDialect.get_table_comment = get_table_comment
 PGDialect._get_column_info = get_column_info  # pylint: disable=protected-access
@@ -307,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
@@ -346,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)

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")

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');
+"""

metadata/ingestion/source/database/query_parser_source.py

@@ -23,10 +23,11 @@ from metadata.ingestion.api.steps import Source
 from metadata.ingestion.connections.test_connections import (
     raise_test_connection_exception,
 )
+from metadata.ingestion.lineage.masker import masked_query_cache
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_test_connection_fn
-from metadata.utils.helpers import get_start_and_end
+from metadata.utils.helpers import get_start_and_end, retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.ssl_manager import get_ssl_connection
 
@@ -49,6 +50,7 @@ class QueryParserSource(Source, ABC):
     database_field: str
     schema_field: str
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,
@@ -64,11 +66,13 @@ class QueryParserSource(Source, ABC):
         self.dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)
         self.source_config = self.config.sourceConfig.config
         self.start, self.end = get_start_and_end(self.source_config.queryLogDuration)
-        self.engine = (
-            get_ssl_connection(self.service_connection) if get_engine else None
-        )
         self.graph = None
 
+        self.engine = None
+        if get_engine:
+            self.engine = get_ssl_connection(self.service_connection)
+            self.test_connection()
+
     @property
     def name(self) -> str:
         return self.service_connection.type.name
@@ -126,9 +130,10 @@ class QueryParserSource(Source, ABC):
         yield self.engine
 
     def close(self):
-        """By default, there is nothing to close"""
+        # Clear the cache
+        masked_query_cache.clear()
 
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(self.engine)
+        result = test_connection_fn(self.metadata, self.engine, self.service_connection)
         raise_test_connection_exception(result)