Ashhar commited on
Commit
a2e9487
·
1 Parent(s): 6d149f9

restrict non-readonly queries

Browse files
Files changed (1) hide show
  1. app.py +27 -6
app.py CHANGED
@@ -239,10 +239,30 @@ def clean_sql_response(response: str) -> str:
239
  return response.strip()
240
 
241
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
242
  def execute_query(query):
243
  if not st.session_state.engine:
244
  return None
245
 
 
 
 
 
 
246
  try:
247
  start_time = time.time()
248
  with st.spinner("Executing SQL query..."):
@@ -299,12 +319,13 @@ def generate_sql_query(user_query):
299
  {chr(10).join(tables_context)}
300
 
301
  Important:
302
- 1. Only return the SQL query, nothing else
303
- 2. The query should be valid PostgreSQL syntax
304
- 3. Do not include any explanations or comments
305
- 4. Make sure to handle NULL values appropriately
306
- 5. If joining tables, use appropriate join conditions based on the schema
307
- 6. Use table names with appropriate qualifiers to avoid ambiguity
 
308
 
309
  User Query: {user_query}
310
  """
 
239
  return response.strip()
240
 
241
 
242
+ def is_read_only_query(query: str) -> bool:
243
+ """Check if the query is read-only (SELECT only)."""
244
+ # Convert query to uppercase for case-insensitive comparison
245
+ query_upper = query.upper()
246
+
247
+ # List of SQL statements that modify data
248
+ modification_statements = [
249
+ 'INSERT', 'UPDATE', 'DELETE', 'DROP', 'CREATE', 'ALTER', 'TRUNCATE',
250
+ 'REPLACE', 'MERGE', 'UPSERT', 'GRANT', 'REVOKE'
251
+ ]
252
+
253
+ # Check if query starts with any modification statement
254
+ return not any(query_upper.strip().startswith(stmt) for stmt in modification_statements)
255
+
256
+
257
  def execute_query(query):
258
  if not st.session_state.engine:
259
  return None
260
 
261
+ # Check if the query is read-only
262
+ if not is_read_only_query(query):
263
+ st.error("Error: Only SELECT queries are allowed for security reasons.")
264
+ return None
265
+
266
  try:
267
  start_time = time.time()
268
  with st.spinner("Executing SQL query..."):
 
319
  {chr(10).join(tables_context)}
320
 
321
  Important:
322
+ 1. Only generate SELECT queries - no INSERT, UPDATE, DELETE, or other data modification statements
323
+ 2. Only return the SQL query, nothing else
324
+ 3. The query should be valid PostgreSQL syntax
325
+ 4. Do not include any explanations or comments
326
+ 5. Make sure to handle NULL values appropriately
327
+ 6. If joining tables, use appropriate join conditions based on the schema
328
+ 7. Use table names with appropriate qualifiers to avoid ambiguity
329
 
330
  User Query: {user_query}
331
  """