Fix #17963: Fix PinotDB Ingestion (#18266)

* Fix #17963: Fix PinotDB Ingestion

* fix conn args

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

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"