From 7c2eb5e09ce6ddb7e10c8c79f6c73bf98cb443ae Mon Sep 17 00:00:00 2001 From: Jordi Mas Date: Thu, 8 May 2025 19:22:16 +0200 Subject: [PATCH] fix(graph): enable setting up the timeout for queries Fixes #170 --- src/langchain_google_spanner/graph_store.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/langchain_google_spanner/graph_store.py b/src/langchain_google_spanner/graph_store.py index a4d857e..043a452 100644 --- a/src/langchain_google_spanner/graph_store.py +++ b/src/langchain_google_spanner/graph_store.py @@ -1197,6 +1197,7 @@ def __init__( instance_id: str, database_id: str, client: Optional[spanner.Client] = None, + timeout: Optional[float] = None, ): """Initializes the Spanner implementation. @@ -1208,11 +1209,14 @@ def __init__( self.client = client or spanner.Client() self.instance = self.client.instance(instance_id) self.database = self.instance.database(database_id) + self.timeout = timeout def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]: param_types = {k: TypeUtility.value_to_param_type(v) for k, v in params.items()} with self.database.snapshot() as snapshot: - rows = snapshot.execute_sql(query, params=params, param_types=param_types) + rows = snapshot.execute_sql( + query, params=params, param_types=param_types, timeout=self.timeout + ) return [ { column: value @@ -1253,6 +1257,7 @@ def __init__( static_node_properties: List[str] = [], static_edge_properties: List[str] = [], impl: Optional[SpannerInterface] = None, + timeout: Optional[float] = None, ): """Initializes SpannerGraphStore. @@ -1267,11 +1272,13 @@ def __init__( properties as static; static_edge_properties: in flexible schema, treat these edge properties as static. + timeout (Optional[float]): The timeout for queries in seconds. """ self.impl = impl or SpannerImpl( instance_id, database_id, client_with_user_agent(client, USER_AGENT_GRAPH_STORE), + timeout=timeout, ) self.schema = SpannerGraphSchema( graph_name,