GraphServer: Remove reliance on CloudSpanner's TypeCode (#48)

* Standardize return type for execute_query() functions

- Also removes cloud-spanner specific fields from the return type.
  Specifically `StructType.Field` is removed from the return type.
  Removing this tightly coupled logic is required to allow new DB
  implementations.

* Abstract away "database" implementations and remove strong coupling of DB implementation with APIs

1. Abstracts SpannerDatabase with clear APIs
2. Introduces CloudSpannerDatabase as an implementation of SpannerDatabase
3. Removes further tight coupling with the cloud spanner client by adding a
   SpannerFieldInfo dataclass to replace usage of StructType.Field

* Introduce exec_env.py to maintain global state + minor bug fix

1. The global database_instances is moved to exec_env.py to avoid
   circular imports.

2. SpannerFiledInfo.typename populated with the correct name now

3. Remove all cloud spanner refs from database.py

* GraphServer: Remove reliance on CloudSpanner's TypeCode

The google.cloud.spanner.TypeCode was used to validate supported
types for properties. However, the type information is always
received as a string from the JS client. It is only internally converted
to a TypeCode for verification.

Comparing against the TypeCode does not make the validation any more
robust than just confirming if the strings themselves are valid type
strings.

Removing reliance on TypeCode also makes GraphServer database agnostic.

* SpannerQueryResult: Change 'error' to 'err'

All consumers of SpannerQueryResult expect `err` instead of `error`.
Without this fix, failures are silent. With this, they're displayed
within the notebook.

Thanks to @cqian23 for pointing this out

---------

Co-authored-by: Sailesh Mukil <mukil.sailesh@gmail.com>
Co-authored-by: Sailesh Mukil Gangatharan <saileshmukil@google.com>

Author: 3 people

spanner_graphs/cloud_database.py

@@ -0,0 +1,153 @@
+# Copyright 2024 Google LLC
+
+# Licensed under the Apache License, Version 2.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://www.apache.org/licenses/LICENSE-2.0
+
+# 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.
+
+"""
+This module contains the cloud-specific implementation for talking to a Spanner database.
+"""
+
+from __future__ import annotations
+import json
+from typing import Any, Dict, List, Tuple
+
+from google.cloud import spanner
+from google.cloud.spanner_v1 import JsonObject
+from google.api_core.client_options import ClientOptions
+from google.cloud.spanner_v1.types import StructType, Type, TypeCode
+import pydata_google_auth
+
+from spanner_graphs.database import SpannerDatabase, MockSpannerDatabase, SpannerQueryResult, SpannerFieldInfo
+
+def _get_default_credentials_with_project():
+    return pydata_google_auth.default(
+        scopes=["https://www.googleapis.com/auth/cloud-platform"], use_local_webserver=False)
+
+def get_as_field_info_list(fields: List[StructType.Field]) -> List[SpannerFieldInfo]:
+  """Converts a list of StructType.Field to a list of SpannerFieldInfo."""
+  return [SpannerFieldInfo(name=field.name, typename=TypeCode(field.type_.code).name) for field in fields]
+
+class CloudSpannerDatabase(SpannerDatabase):
+    """Concrete implementation for Spanner database on the cloud."""
+    def __init__(self, project_id: str, instance_id: str,
+                 database_id: str) -> None:
+        credentials, _ = _get_default_credentials_with_project()
+        self.client = spanner.Client(
+            project=project_id, credentials=credentials, client_options=ClientOptions(quota_project_id=project_id))
+        self.instance = self.client.instance(instance_id)
+        self.database = self.instance.database(database_id)
+        self.schema_json: Any | None = None
+
+    def __repr__(self) -> str:
+        return (f"<CloudSpannerDatabase["
+                f"project:{self.client.project_name},"
+                f"instance:{self.instance.name},"
+                f"db:{self.database.name}]>")
+
+    def _extract_graph_name(self, query: str) -> str:
+        words = query.strip().split()
+        if len(words) < 3:
+            raise ValueError("invalid query: must contain at least (GRAPH, graph_name and query)")
+
+        if words[0].upper() != "GRAPH":
+            raise ValueError("invalid query: GRAPH must be the first word")
+
+        return words[1]
+
+    def _get_schema_for_graph(self, graph_query: str) -> Any | None:
+        try:
+            graph_name = self._extract_graph_name(graph_query)
+        except ValueError:
+            return None
+
+        with self.database.snapshot() as snapshot:
+            schema_query = """
+            SELECT property_graph_name, property_graph_metadata_json
+            FROM information_schema.property_graphs
+            WHERE property_graph_name = @graph_name
+            """
+            params = {"graph_name": graph_name}
+            param_type = {"graph_name": spanner.param_types.STRING}
+
+            result = snapshot.execute_sql(schema_query, params=params, param_types=param_type)
+            schema_rows = list(result)
+
+            if schema_rows:
+                return schema_rows[0][1]
+            else:
+                return None
+
+    def execute_query(
+        self,
+        query: str,
+        limit: int = None,
+        is_test_query: bool = False,
+    ) -> SpannerQueryResult:
+        """
+        This method executes the provided `query`
+
+        Args:
+            query: The SQL query to execute against the database
+            limit: An optional limit for the number of rows to return
+            is_test_query: If true, skips schema fetching for graph queries.
+
+        Returns:
+            A `SpannerQueryResult`
+        """
+        self.schema_json = None
+        if not is_test_query:
+            self.schema_json = self._get_schema_for_graph(query)
+
+        with self.database.snapshot() as snapshot:
+            params = None
+            param_types = None
+            if limit and limit > 0:
+                params = dict(limit=limit)
+
+            try:
+                results = snapshot.execute_sql(query, params=params, param_types=param_types)
+                rows = list(results)
+            except Exception as e:
+                return SpannerQueryResult(
+                    data={},
+                    fields=[],
+                    rows=[],
+                    schema_json=self.schema_json,
+                    err=e
+                )
+
+            fields: List[SpannerFieldInfo] = get_as_field_info_list(results.fields)
+            data = {field.name: [] for field in fields}
+
+            if len(fields) == 0:
+                return SpannerQueryResult(
+                    data=data,
+                    fields=fields,
+                    rows=rows,
+                    schema_json=self.schema_json,
+                    err=None
+                )
+
+            for row_data in rows:
+                for field, value in zip(fields, row_data):
+                    if isinstance(value, JsonObject):
+                        data[field.name].append(json.loads(value.serialize()))
+                    else:
+                        data[field.name].append(value)
+
+            return SpannerQueryResult(
+                data=data,
+                fields=fields,
+                rows=rows,
+                schema_json=self.schema_json,
+                err=None
+            )

spanner_graphs/conversion.py

@@ -21,12 +21,11 @@
 from typing import Any, List, Dict, Tuple
 import json
 
-from google.cloud.spanner_v1.types import TypeCode, StructType
-
+from spanner_graphs.database import SpannerFieldInfo
 from spanner_graphs.graph_entities import Node, Edge
 from spanner_graphs.schema_manager import SchemaManager
 
-def get_nodes_edges(data: Dict[str, List[Any]], fields: List[StructType.Field], schema_json: dict = None) -> Tuple[List[Node], List[Edge]]:
+def get_nodes_edges(data: Dict[str, List[Any]], fields: List[SpannerFieldInfo], schema_json: dict = None) -> Tuple[List[Node], List[Edge]]:
     schema_manager = SchemaManager(schema_json)
     nodes: List[Node] = []
     edges: List[Edge] = []
@@ -37,15 +36,15 @@ def get_nodes_edges(data: Dict[str, List[Any]], fields: List[StructType.Field],
     for field in fields:
         column_name = field.name
         column_data = data[column_name]
-        
+
         # Only process JSON and Array of JSON types
-        if field.type_.code not in [TypeCode.JSON, TypeCode.ARRAY]:
+        if field.typename not in ["JSON", "ARRAY"]:
             continue
 
         # Process each value in the column
         for value in column_data:
             items_to_process = []
-            
+
             # Handle both single JSON and arrays of JSON
             if isinstance(value, list):
                 items_to_process.extend(value)
@@ -92,4 +91,4 @@ def get_nodes_edges(data: Dict[str, List[Any]], fields: List[StructType.Field],
         nodes.append(Node.make_intermediate(identifier))
         node_identifiers.add(identifier)
 
-    return nodes, edges
+    return nodes, edges