Upload environment.py

server/environment.py

@@ -210,31 +210,71 @@ class SQLQualityEnvironment:
         """
         Handle a QUERY action.
 
-        Executes the SQL on the episode connection.
+        Executes the SQL on the episode connection (supports multi-statement).
         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()
 
-        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]
+        # Split multiple statements by semicolon, but preserve semicolons in strings
+        statements = []
+        current = ""
+        in_string = False
+        string_char = None
+        i = 0
+        while i < len(sql):
+            char = sql[i]
+            if char in ("'", '"') and (i == 0 or sql[i-1] != "\\"):
+                if not in_string:
+                    in_string = True
+                    string_char = char
+                elif char == string_char:
+                    in_string = False
+                    string_char = None
+            if char == ";" and not in_string:
+                stmt = current.strip()
+                if stmt:
+                    statements.append(stmt)
+                current = ""
             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
+                current += char
+            i += 1
+        # Add final statement
+        if current.strip():
+            statements.append(current.strip())
+
+        # If no statements found, treat as single statement
+        if not statements:
+            statements = [sql] if sql else []
+
+        result_preview = ""
+        step_reward = 0.0
+        last_error = None
 
+        for stmt in statements:
+            if not stmt:
+                continue
+            try:
+                cursor = self._original_conn.execute(stmt)
+                # 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]
+                else:
+                    # For INSERT/UPDATE/DELETE, show affected rows
+                    result_preview = f"Query executed successfully. Rows affected: {cursor.rowcount}"
+            except sqlite3.Error as e:
+                last_error = str(e)
+                result_preview = f"SQL Error: {e}"
+                break  # Stop on first error
+
+        # Small reward for targeted queries (only check first statement for relevance)
+        if statements and not last_error:
+            from server.graders import _query_reveals_issue
             step_reward = (
-                0.05 if _query_reveals_issue(sql, self._ground_truth) else 0.01
+                0.05 if _query_reveals_issue(statements[0], 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(