Add schema related to get_table_names and get_columns

Author: splasky

sqlalchemy_datastore/_helpers.py

@@ -20,7 +20,7 @@
 import functools
 import re
 import os
-from typing import Optional
+from typing import Optional, Tuple
 
 from google.api_core import client_info
 import google.auth
@@ -58,7 +58,8 @@ def create_datastore_client(
     credentials_base64: Optional[str] = None,
     project_id: Optional[str] = None,
     user_agent: Optional[client_info.ClientInfo] = None,
-) -> datastore.Client:
+    database: Optional[str] = None
+) -> Tuple[datastore.Client, service_account.Credentials]:
     """Construct a BigQuery client object.
 
     Args:
@@ -77,7 +78,7 @@ def create_datastore_client(
     """
 
     default_project = None
-
+    database = database if database != "(default)" else None
     if os.getenv("DATASTORE_EMULATOR_HOST") is not None:
         client = datastore.Client(project=project_id)
         return client
@@ -109,7 +110,8 @@ def create_datastore_client(
         client_info=info,
         project=project_id,
         credentials=credentials,
-    )
+        database=database
+    ), credentials
 
 
 def substitute_re_method(r, flags=0, repl=None):

sqlalchemy_datastore/base.py

@@ -18,7 +18,7 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 import os
 import logging
-from typing import Any, List
+from typing import Any, List, Optional, Dict
 from concurrent import futures
 
 from . import _types
@@ -32,9 +32,8 @@
 from sqlalchemy.sql.expression import TextClause
 
 from google.cloud import firestore_admin_v1
-from google.cloud.firestore_admin_v1.types import ListDatabasesResponse
-from google.api_core.exceptions import GoogleAPIError
-from google.api_core import client_info
+from google.cloud.firestore_admin_v1.types import ListDatabasesResponse, Database
+from google.oauth2 import service_account
 
 logger = logging.getLogger('sqlalchemy.dialects.CloudDatastore')
 
@@ -354,6 +353,7 @@ def __init__(
         self.dataset_id = None
         self.list_tables_page_size = list_tables_page_size
         self._client = None
+        self.credentials = None
 
     @classmethod
     def dbapi(cls):
@@ -363,7 +363,8 @@ def dbapi(cls):
     def do_ping(self, dbapi_connection):
         """Performs a simple operation to check if the connection is still alive."""
         try:
-            # Basic connectivity check
+            query = self._client.query(kind="__kind__")
+            query.fetch(limit=1, timeout=30)
             return True
         except Exception:
             return False
@@ -406,11 +407,12 @@ def create_connect_args(self, url):
         if user_supplied_client:
             return ([], {})
         else:
-            client = create_datastore_client(
+            client, credentials = create_datastore_client(
                 credentials_path=self.credentials_path,
                 credentials_info=self.credentials_info,
                 credentials_base64=self.credentials_base64,
                 project_id=self.billing_project_id,
+                database=None,
             )
             self.project_id = self.project_id if self.project_id else client.project
             self.billing_project_id = (
@@ -423,13 +425,14 @@ def create_connect_args(self, url):
             )
 
         self._client = client
+        self.credentials = credentials
         setattr(self._client, "credentials_path", self.credentials_path)
         setattr(self._client, "credentials_info", self.credentials_info)
         setattr(self._client, "credentials_base64", self.credentials_base64)
         return ([], {"client": client})
 
-    def get_schema_names(self, connection, **kw):
-        return self._list_datastore_databases(self.credentials_info, self.project_id)
+    def get_schema_names(self, connection: Connection, **kw) -> Optional[List[str]]:
+        return self._list_datastore_databases(self.credentials, self.project_id)
 
     def _list_datastore_databases(self, cred: service_account.Credentials, project_id: str) -> Optional[List[str]]:
         """Lists all Datastore databases for a given Google Cloud project.
@@ -454,8 +457,13 @@ def get_database_short_name(database: Database) -> Optional[List[str]]:
             logging.error(e)
         return []
 
-    def get_table_names(self, connection, schema: str | None = None, **kw):
-        client = self._client
+    def get_table_names(self, connection: Connection, schema: str | None = None, **kw) -> Optional[List[str]]:
+        client, _ = create_datastore_client(
+            credentials_path=self.credentials_path,
+            credentials_info=self.credentials_info,
+            credentials_base64=self.credentials_base64,
+            database=schema
+        )
         query = client.query(kind="__kind__")
         kinds = list(query.fetch())
 
@@ -469,7 +477,12 @@ def get_kind_name(kind):
 
     def get_columns(self, connection: Connection, table_name: str, schema: str | None = None, **kw):
         """Retrieve column information from the database with optimized parallel processing."""
-        client = self._client
+        client, _ = create_datastore_client(
+            credentials_path=self.credentials_path,
+            credentials_info=self.credentials_info,
+            credentials_base64=self.credentials_base64,
+            database=schema
+        )
         ancestor_key = client.key("__kind__", table_name)
         query = client.query(kind="__property__", ancestor=ancestor_key)
         properties = list(query.fetch())