Ashhar
commited on
Commit
·
a2e9487
1
Parent(s):
6d149f9
restrict non-readonly queries
Browse files
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
|
303 |
-
2.
|
304 |
-
3.
|
305 |
-
4.
|
306 |
-
5.
|
307 |
-
6.
|
|
|
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 |
"""
|