File tree 2 files changed +13
-1
lines changed
libs/knowledge-store/ragstack_knowledge_store
2 files changed +13
-1
lines changed Original file line number Diff line number Diff line change @@ -67,6 +67,7 @@ def execute(
67
67
query : PreparedStatement ,
68
68
parameters : tuple [Any , ...] | None = None ,
69
69
callback : _Callback | None = None ,
70
+ timeout : float | None = None ,
70
71
) -> None :
71
72
"""Execute a query concurrently.
72
73
@@ -77,6 +78,7 @@ def execute(
77
78
query: The query to execute.
78
79
parameters: Parameter tuple for the query. Defaults to `None`.
79
80
callback: Callback to apply to the results. Defaults to `None`.
81
+ timeout: Timeout to use (if not the session default).
80
82
"""
81
83
# TODO: We could have some form of throttling, where we track the number
82
84
# of pending calls and queue things if it exceed some threshold.
@@ -86,7 +88,14 @@ def execute(
86
88
if self ._error is not None :
87
89
return
88
90
89
- future : ResponseFuture = self ._session .execute_async (query , parameters )
91
+ execute_kwargs = {}
92
+ if timeout is not None :
93
+ execute_kwargs ["timeout" ] = timeout
94
+ future : ResponseFuture = self ._session .execute_async (
95
+ query ,
96
+ parameters ,
97
+ ** execute_kwargs ,
98
+ )
90
99
future .add_callbacks (
91
100
self ._handle_result ,
92
101
self ._handle_error ,
Original file line number Diff line number Diff line change @@ -157,7 +157,9 @@ def __init__(
157
157
keyspace : str | None = None ,
158
158
setup_mode : SetupMode = SetupMode .SYNC ,
159
159
metadata_indexing : MetadataIndexingType = "all" ,
160
+ insert_timeout : float = 30.0 ,
160
161
):
162
+ self ._insert_timeout = insert_timeout
161
163
if targets_table :
162
164
logger .warning (
163
165
"The 'targets_table' parameter is deprecated "
@@ -318,6 +320,7 @@ def add_nodes(
318
320
metadata_blob ,
319
321
metadata_s ,
320
322
),
323
+ timeout = self ._insert_timeout ,
321
324
)
322
325
323
326
return node_ids
You can’t perform that action at this time.
0 commit comments