diff --git a/ingestion/src/metadata/ingestion/source/database/pinotdb/connection.py b/ingestion/src/metadata/ingestion/source/database/pinotdb/connection.py index 2f25281d670..63d25adcbd6 100644 --- a/ingestion/src/metadata/ingestion/source/database/pinotdb/connection.py +++ b/ingestion/src/metadata/ingestion/source/database/pinotdb/connection.py @@ -12,6 +12,7 @@ """ Source connection handler """ +from copy import deepcopy from typing import Optional from sqlalchemy.engine import Engine @@ -26,6 +27,7 @@ from metadata.ingestion.connections.builders import ( create_generic_db_connection, get_connection_args_common, get_connection_url_common, + init_empty_connection_arguments, ) from metadata.ingestion.connections.test_connections import test_connection_db_common from metadata.ingestion.ometa.ometa_api import OpenMetadata @@ -41,8 +43,16 @@ def get_connection(connection: PinotDBConnection) -> Engine: """ Create connection """ + # TODO: Rename database field to DatabaseSchema + # Pinot does not support multi database concept + if connection.database is not None: + if not connection.connectionArguments: + connection.connectionArguments = init_empty_connection_arguments() + connection.connectionArguments.root["database"] = connection.database + connection_copy = deepcopy(connection) + connection_copy.database = None return create_generic_db_connection( - connection=connection, + connection=connection_copy, get_connection_url_fn=get_connection_url, get_connection_args_fn=get_connection_args_common, ) diff --git a/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py b/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py index fed71547098..f38b3d82ba1 100644 --- a/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/pinotdb/metadata.py @@ -9,7 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """PinotDb source module""" -from typing import Optional +from typing import Iterable, Optional from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import ( PinotDBConnection, @@ -39,3 +39,15 @@ class PinotdbSource(CommonDbSourceService): f"Expected PinotdbConnection, but got {connection}" ) return cls(config, metadata) + + def get_database_names(self) -> Iterable[str]: + """ + Default case with a single database. + + It might come informed - or not - from the source. + + Sources with multiple databases should overwrite this and + apply the necessary filters. + """ + # TODO: Add databaseDisplayName field in PinotDBConnection + yield "default"