diff --git a/src/langchain_google_spanner/graph_store.py b/src/langchain_google_spanner/graph_store.py index b8cefc05..2236ba22 100644 --- a/src/langchain_google_spanner/graph_store.py +++ b/src/langchain_google_spanner/graph_store.py @@ -1230,6 +1230,7 @@ def __init__( instance_id: str, database_id: str, client: Optional[spanner.Client] = None, + timeout: Optional[float] = None, ): """Initializes the Spanner implementation. @@ -1241,11 +1242,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 @@ -1286,6 +1290,7 @@ def __init__( static_node_properties: List[str] = [], static_edge_properties: List[str] = [], impl: Optional[SpannerInterface] = None, + timeout: Optional[float] = None, ): """Initializes SpannerGraphStore. @@ -1300,11 +1305,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,