Skip to content

Commit 177b9af

Browse files
dihenkerdishaprakash
dihenker
and
dishaprakash
authored
Refactor: parameterize SQL query values (#212)
* Refactor: parameterize SQL query values * Format with `black` linter * Fix variable assignment typing --------- Co-authored-by: dishaprakash <57954147+dishaprakash@users.noreply.github.com>
1 parent f8a230c commit 177b9af

4 files changed

Lines changed: 66 additions & 43 deletions

File tree

‎src/llama_index_cloud_sql_pg/async_chat_store.py‎

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,9 @@ async def __aexecute_query(self, query, params=None):
104104
await conn.execute(text(query), params)
105105
await conn.commit()
106106

107-
async def __afetch_query(self, query):
107+
async def __afetch_query(self, query, params=None):
108108
async with self._engine.connect() as conn:
109-
result = await conn.execute(text(query))
109+
result = await conn.execute(text(query), params)
110110
result_map = result.mappings()
111111
results = result_map.fetchall()
112112
await conn.commit()
@@ -128,8 +128,9 @@ async def aset_messages(self, key: str, messages: List[ChatMessage]) -> None:
128128
None
129129
130130
"""
131-
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE key = '{key}'; """
132-
await self.__aexecute_query(query)
131+
delete_query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE key = :key; """
132+
delete_params = {"key": key}
133+
await self.__aexecute_query(delete_query, delete_params)
133134
insert_query = f"""
134135
INSERT INTO "{self._schema_name}"."{self._table_name}" (key, message)
135136
VALUES (:key, :message);"""
@@ -154,8 +155,9 @@ async def aget_messages(self, key: str) -> List[ChatMessage]:
154155
List[ChatMessage]: A list of `ChatMessage` objects associated with the provided key.
155156
If no messages are found, an empty list is returned.
156157
"""
157-
query = f"""SELECT message from "{self._schema_name}"."{self._table_name}" WHERE key = '{key}' ORDER BY id;"""
158-
results = await self.__afetch_query(query)
158+
query = f"""SELECT message from "{self._schema_name}"."{self._table_name}" WHERE key = :key ORDER BY id;"""
159+
params = {"key": key}
160+
results = await self.__afetch_query(query, params)
159161
if results:
160162
return [
161163
ChatMessage.model_validate(result.get("message")) for result in results
@@ -189,8 +191,9 @@ async def adelete_messages(self, key: str) -> Optional[List[ChatMessage]]:
189191
Optional[List[ChatMessage]]: A list of `ChatMessage` objects that were deleted, or `None` if no messages
190192
were associated with the key or could be deleted.
191193
"""
192-
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE key = '{key}' RETURNING *; """
193-
results = await self.__afetch_query(query)
194+
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE key = :key RETURNING *; """
195+
params = {"key": key}
196+
results = await self.__afetch_query(query, params)
194197
if results:
195198
return [
196199
ChatMessage.model_validate(result.get("message")) for result in results
@@ -208,14 +211,16 @@ async def adelete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
208211
Optional[ChatMessage]: The `ChatMessage` object that was deleted, or `None` if no message
209212
was associated with the key or could be deleted.
210213
"""
211-
query = f"""SELECT * from "{self._schema_name}"."{self._table_name}" WHERE key = '{key}' ORDER BY id;"""
212-
results = await self.__afetch_query(query)
214+
query = f"""SELECT * from "{self._schema_name}"."{self._table_name}" WHERE key = :key ORDER BY id;"""
215+
params = {"key": key}
216+
results = await self.__afetch_query(query, params)
213217
if results:
214218
if idx >= len(results):
215219
return None
216220
id_to_be_deleted = results[idx].get("id")
217-
delete_query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE id = '{id_to_be_deleted}' RETURNING *;"""
218-
result = await self.__afetch_query(delete_query)
221+
delete_query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE id = :id_to_be_deleted RETURNING *;"""
222+
delete_params = {"id_to_be_deleted": id_to_be_deleted}
223+
result = await self.__afetch_query(delete_query, delete_params)
219224
result = result[0]
220225
if result:
221226
return ChatMessage.model_validate(result.get("message"))
@@ -232,12 +237,14 @@ async def adelete_last_message(self, key: str) -> Optional[ChatMessage]:
232237
Optional[ChatMessage]: The `ChatMessage` object that was deleted, or `None` if no message
233238
was associated with the key or could be deleted.
234239
"""
235-
query = f"""SELECT * from "{self._schema_name}"."{self._table_name}" WHERE key = '{key}' ORDER BY id DESC LIMIT 1;"""
236-
results = await self.__afetch_query(query)
240+
query = f"""SELECT * from "{self._schema_name}"."{self._table_name}" WHERE key = :key ORDER BY id DESC LIMIT 1;"""
241+
params = {"key": key}
242+
results = await self.__afetch_query(query, params)
237243
if results:
238244
id_to_be_deleted = results[0].get("id")
239-
delete_query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE id = '{id_to_be_deleted}' RETURNING *;"""
240-
result = await self.__afetch_query(delete_query)
245+
delete_query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE id = :id_to_be_deleted RETURNING *;"""
246+
delete_params = {"id_to_be_deleted": id_to_be_deleted}
247+
result = await self.__afetch_query(delete_query, delete_params)
241248
result = result[0]
242249
if result:
243250
return ChatMessage.model_validate(result.get("message"))

‎src/llama_index_cloud_sql_pg/async_document_store.py‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -110,9 +110,9 @@ async def __aexecute_query(self, query, params):
110110
await conn.commit()
111111
return None
112112

113-
async def __afetch_query(self, query):
113+
async def __afetch_query(self, query, params=None):
114114
async with self._engine.connect() as conn:
115-
result = await conn.execute(text(query))
115+
result = await conn.execute(text(query), params)
116116
result_map = result.mappings()
117117
results = result_map.fetchall()
118118
await conn.commit()
@@ -159,8 +159,9 @@ async def _delete_from_table(self, id: str) -> Sequence[RowMapping]:
159159
Returns:
160160
List of deleted rows.
161161
"""
162-
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE id = '{id}' RETURNING *; """
163-
result = await self.__afetch_query(query)
162+
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE id = :id RETURNING *; """
163+
params = {"id": id}
164+
result = await self.__afetch_query(query, params)
164165
return result
165166

166167
async def async_add_documents(
@@ -254,8 +255,9 @@ async def aget_document(
254255
Returns:
255256
Optional[BaseNode]: Returns a `BaseNode` object if the document is found
256257
"""
257-
query = f"""SELECT node_data from "{self._schema_name}"."{self._table_name}" WHERE id = '{doc_id}';"""
258-
result = await self.__afetch_query(query)
258+
query = f"""SELECT node_data from "{self._schema_name}"."{self._table_name}" WHERE id = :doc_id;"""
259+
params = {"doc_id": doc_id}
260+
result = await self.__afetch_query(query, params)
259261

260262
if result:
261263
result = result[0]
@@ -276,9 +278,10 @@ async def aget_ref_doc_info(self, ref_doc_id: str) -> Optional[RefDocInfo]:
276278
Returns:
277279
Optional[RefDocInfo]: Returns a `RefDocInfo` object if it exists.
278280
"""
279-
query = f"""select id, node_data from "{self._schema_name}"."{self._table_name}" where ref_doc_id = '{ref_doc_id}'"""
281+
query = f"""select id, node_data from "{self._schema_name}"."{self._table_name}" where ref_doc_id = :ref_doc_id"""
282+
params = {"ref_doc_id": ref_doc_id}
280283

281-
rows = await self.__afetch_query(query)
284+
rows = await self.__afetch_query(query, params)
282285
node_ids = []
283286
merged_metadata = {}
284287

@@ -350,8 +353,9 @@ async def adocument_exists(self, doc_id: str) -> bool:
350353
Returns:
351354
bool : True if document exists in the table.
352355
"""
353-
query = f"""SELECT id from "{self._schema_name}"."{self._table_name}" WHERE id = '{doc_id}' LIMIT 1;"""
354-
result = await self.__afetch_query(query)
356+
query = f"""SELECT id from "{self._schema_name}"."{self._table_name}" WHERE id = :doc_id LIMIT 1;"""
357+
params = {"doc_id": doc_id}
358+
result = await self.__afetch_query(query, params)
355359
return bool(result)
356360

357361
async def _get_ref_doc_child_node_ids(
@@ -366,8 +370,9 @@ async def _get_ref_doc_child_node_ids(
366370
list # List of all nodes that refer to ref_doc_id
367371
]
368372
]"""
369-
query = f"""select id from "{self._schema_name}"."{self._table_name}" where ref_doc_id = '{ref_doc_id}';"""
370-
results = await self.__afetch_query(query)
373+
query = f"""select id from "{self._schema_name}"."{self._table_name}" where ref_doc_id = :ref_doc_id;"""
374+
params = {"ref_doc_id": ref_doc_id}
375+
results = await self.__afetch_query(query, params)
371376
result = {"node_ids": [item["id"] for item in results]}
372377
return result
373378

@@ -423,7 +428,8 @@ async def adelete_ref_doc(self, ref_doc_id: str, raise_error: bool = True) -> No
423428
return
424429

425430
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE ref_doc_id = :ref_doc_id;"""
426-
await self.__aexecute_query(query, {"ref_doc_id": ref_doc_id})
431+
params = {"ref_doc_id": ref_doc_id}
432+
await self.__aexecute_query(query, params)
427433

428434
await self._delete_from_table(ref_doc_id)
429435

@@ -465,8 +471,9 @@ async def aget_document_hash(self, doc_id: str) -> Optional[str]:
465471
str # hash for the given doc_id
466472
]
467473
"""
468-
query = f"""SELECT id, doc_hash from "{self._schema_name}"."{self._table_name}" WHERE id = '{doc_id}' LIMIT 1;"""
469-
row = await self.__afetch_query(query)
474+
query = f"""SELECT id, doc_hash from "{self._schema_name}"."{self._table_name}" WHERE id = :doc_id LIMIT 1;"""
475+
params = {"doc_id": doc_id}
476+
row = await self.__afetch_query(query, params)
470477

471478
if row:
472479
return row[0].get("doc_hash", None)

‎src/llama_index_cloud_sql_pg/async_index_store.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,9 @@ async def __aexecute_query(self, query, params=None):
104104
await conn.execute(text(query), params)
105105
await conn.commit()
106106

107-
async def __afetch_query(self, query):
107+
async def __afetch_query(self, query, params=None):
108108
async with self._engine.connect() as conn:
109-
result = await conn.execute(text(query))
109+
result = await conn.execute(text(query), params)
110110
result_map = result.mappings()
111111
results = result_map.fetchall()
112112
await conn.commit()
@@ -171,8 +171,9 @@ async def adelete_index_struct(self, key: str) -> None:
171171
key (str): index struct key
172172
173173
"""
174-
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE index_id = '{key}'; """
175-
await self.__aexecute_query(query)
174+
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE index_id = :key; """
175+
params = {"key": key}
176+
await self.__aexecute_query(query, params)
176177
return
177178

178179
async def aget_index_struct(
@@ -191,8 +192,9 @@ async def aget_index_struct(
191192
warnings.warn("No struct_id specified and more than one struct exists.")
192193
return None
193194
else:
194-
query = f"""SELECT * from "{self._schema_name}"."{self._table_name}" WHERE index_id = '{struct_id}';"""
195-
result = await self.__afetch_query(query)
195+
query = f"""SELECT * from "{self._schema_name}"."{self._table_name}" WHERE index_id = :struct_id;"""
196+
params = {"struct_id": struct_id}
197+
result = await self.__afetch_query(query, params)
196198
if result:
197199
json = result[0]
198200
if json is None:

‎src/llama_index_cloud_sql_pg/async_vector_store.py‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -155,9 +155,10 @@ async def create(
155155
Returns:
156156
AsyncPostgresVectorStore
157157
"""
158-
stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}' AND table_schema = '{schema_name}'"
158+
stmt = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = :table_name AND table_schema = :schema_name"
159+
params = {"table_name": table_name, "schema_name": schema_name}
159160
async with engine._pool.connect() as conn:
160-
result = await conn.execute(text(stmt))
161+
result = await conn.execute(text(stmt), params)
161162
result_map = result.mappings()
162163
results = result_map.fetchall()
163164
columns = {}
@@ -280,9 +281,10 @@ async def async_add(self, nodes: Sequence[BaseNode], **kwargs: Any) -> list[str]
280281

281282
async def adelete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
282283
"""Asynchronously delete nodes belonging to provided parent document from the table."""
283-
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE {self._ref_doc_id_column} = '{ref_doc_id}'"""
284+
query = f"""DELETE FROM "{self._schema_name}"."{self._table_name}" WHERE {self._ref_doc_id_column} = :ref_doc_id;"""
285+
params = {"ref_doc_id": ref_doc_id}
284286
async with self._engine.connect() as conn:
285-
await conn.execute(text(query))
287+
await conn.execute(text(query), params)
286288
await conn.commit()
287289

288290
async def adelete_nodes(
@@ -467,10 +469,15 @@ async def is_valid_index(
467469
query = f"""
468470
SELECT tablename, indexname
469471
FROM pg_indexes
470-
WHERE tablename = '{self._table_name}' AND schemaname = '{self._schema_name}' AND indexname = '{index_name}';
472+
WHERE tablename = :table_name AND schemaname = :schema_name AND indexname = :index_name;
471473
"""
474+
params = {
475+
"table_name": self._table_name,
476+
"schema_name": self._schema_name,
477+
"index_name": index_name,
478+
}
472479
async with self._engine.connect() as conn:
473-
result = await conn.execute(text(query))
480+
result = await conn.execute(text(query), params)
474481
result_map = result.mappings()
475482
results = result_map.fetchall()
476483
return bool(len(results) == 1)

0 commit comments

Comments
 (0)