Skip to content

Commit 1c5cf19

Browse files
feat(graph): enable setting up the timeout for queries (#178)
Fixes #170 Co-authored-by: Averi Kitsch <[email protected]>
1 parent 67b3950 commit 1c5cf19

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

src/langchain_google_spanner/graph_store.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1230,6 +1230,7 @@ def __init__(
12301230
instance_id: str,
12311231
database_id: str,
12321232
client: Optional[spanner.Client] = None,
1233+
timeout: Optional[float] = None,
12331234
):
12341235
"""Initializes the Spanner implementation.
12351236
@@ -1241,11 +1242,14 @@ def __init__(
12411242
self.client = client or spanner.Client()
12421243
self.instance = self.client.instance(instance_id)
12431244
self.database = self.instance.database(database_id)
1245+
self.timeout = timeout
12441246

12451247
def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
12461248
param_types = {k: TypeUtility.value_to_param_type(v) for k, v in params.items()}
12471249
with self.database.snapshot() as snapshot:
1248-
rows = snapshot.execute_sql(query, params=params, param_types=param_types)
1250+
rows = snapshot.execute_sql(
1251+
query, params=params, param_types=param_types, timeout=self.timeout
1252+
)
12491253
return [
12501254
{
12511255
column: value
@@ -1286,6 +1290,7 @@ def __init__(
12861290
static_node_properties: List[str] = [],
12871291
static_edge_properties: List[str] = [],
12881292
impl: Optional[SpannerInterface] = None,
1293+
timeout: Optional[float] = None,
12891294
):
12901295
"""Initializes SpannerGraphStore.
12911296
@@ -1300,11 +1305,13 @@ def __init__(
13001305
properties as static;
13011306
static_edge_properties: in flexible schema, treat these edge
13021307
properties as static.
1308+
timeout (Optional[float]): The timeout for queries in seconds.
13031309
"""
13041310
self.impl = impl or SpannerImpl(
13051311
instance_id,
13061312
database_id,
13071313
client_with_user_agent(client, USER_AGENT_GRAPH_STORE),
1314+
timeout=timeout,
13081315
)
13091316
self.schema = SpannerGraphSchema(
13101317
graph_name,

0 commit comments

Comments
 (0)