Upload environment.py

server/environment.py

@@ -208,33 +208,33 @@ class SQLQualityEnvironment:
 
     def _handle_query(self, action: SQLAction) -> StepResponse:
         """
-        Handle a read-only QUERY action.
+        Handle a QUERY action.
 
-        Executes the SQL on the episode connection (read-only SELECT enforced).
+        Executes the SQL on the episode connection.
         Returns a result preview in the observation. Does NOT end the episode.
         Awards small step_reward (0.05) for queries that target the right area.
         """
         sql = (action.sql or "").strip()
 
-        # Enforce read-only: only SELECT allowed
-        if not sql.upper().startswith("SELECT"):
-            step_reward = 0.0
-            result_preview = "Error: QUERY action only accepts SELECT statements."
-        else:
-            try:
-                cursor = self._original_conn.execute(sql)
+        try:
+            cursor = self._original_conn.execute(sql)
+            # Check if query returns rows (SELECT) or just modifies (INSERT/UPDATE/DELETE)
+            if cursor.description:
                 rows = cursor.fetchmany(5)
                 cols = [d[0] for d in cursor.description]
                 result_preview = str([dict(zip(cols, row)) for row in rows])[:300]
-                # Small reward for targeted queries
-                from server.graders import _query_reveals_issue
+            else:
+                # For INSERT/UPDATE/DELETE, show affected rows
+                result_preview = f"Query executed successfully. Rows affected: {cursor.rowcount}"
+            # Small reward for targeted queries
+            from server.graders import _query_reveals_issue
 
-                step_reward = (
-                    0.05 if _query_reveals_issue(sql, self._ground_truth) else 0.01
-                )
-            except sqlite3.Error as e:
-                result_preview = f"SQL Error: {e}"
-                step_reward = 0.0
+            step_reward = (
+                0.05 if _query_reveals_issue(sql, self._ground_truth) else 0.01
+            )
+        except sqlite3.Error as e:
+            result_preview = f"SQL Error: {e}"
+            step_reward = 0.0
 
         # Add to trajectory
         self._trajectory.append(