""" ETL Pipeline Fixer — Environment Implementation ================================================ Phase 3: Multi-level task logic. Difficulty levels: - easy: Schema mismatch (fix INSERT column name) - medium: Silent data loss (change INNER JOIN to LEFT JOIN) - hard: Concurrency (fix 'database is locked' via threading.Lock) """ from __future__ import annotations import shutil import sqlite3 import subprocess import sys import textwrap import uuid from pathlib import Path from typing import Any, Literal, Optional from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import State try: from .models import ( EditCodeAction, EtlAction, EtlObservation, EtlReward, QueryDbAction, RunPipelineAction, ) except ImportError: from models import ( EditCodeAction, EtlAction, EtlObservation, EtlReward, QueryDbAction, RunPipelineAction, ) # --------------------------------------------------------------------------- # ── Types ────────────────────────────────────────────────────────────────── # --------------------------------------------------------------------------- DifficultyLevel = Literal["easy", "medium", "hard"] # --------------------------------------------------------------------------- # ── Task 1 (Easy) constants ──────────────────────────────────────────────── # --------------------------------------------------------------------------- _EASY_EXPECTED_ROWS: list[tuple[int, str]] = [ (1, "Alice Smith"), (2, "Bob Jones"), (3, "Charlie Brown"), ] _EASY_SOURCE_CSV: str = textwrap.dedent("""\ user_id,full_name 1,Alice Smith 2,Bob Jones 3,Charlie Brown """) _EASY_PIPELINE: str = textwrap.dedent("""\ \"\"\"ETL pipeline — Task 1: load source_data.csv into the users table.\"\"\" import csv import sqlite3 # --- Extract --- rows = [] with open("source_data.csv", newline="") as f: reader = csv.DictReader(f) for row in reader: rows.append(row) # --- Load --- conn = sqlite3.connect("etl_target.db") cursor = conn.cursor() # BUG: 'user_id' does not exist in the target schema — should be 'id'. for row in rows: cursor.execute( "INSERT INTO users (user_id, name) VALUES (?, ?)", (row["user_id"], row["full_name"]), ) conn.commit() conn.close() print(f"Loaded {len(rows)} rows into users table.") """) # --------------------------------------------------------------------------- # ── Task 2 (Medium) constants ────────────────────────────────────────────── # --------------------------------------------------------------------------- _MEDIUM_PIPELINE: str = textwrap.dedent("""\ \"\"\"ETL pipeline — Task 2: Create user_summary from users and purchases.\"\"\" import sqlite3 conn = sqlite3.connect("etl_target.db") cursor = conn.cursor() # BUG: INNER JOIN drops users without purchases. # The fix: change to LEFT JOIN to retain all 5 users. query = ''' CREATE TABLE user_summary AS SELECT u.id, u.name, COUNT(p.id) as purchase_count FROM users u INNER JOIN purchases p ON u.id = p.user_id GROUP BY u.id, u.name ''' cursor.execute("DROP TABLE IF EXISTS user_summary") cursor.execute(query) conn.commit() conn.close() print("User summary created.") """) # --------------------------------------------------------------------------- # ── Task 3 (Hard) constants ──────────────────────────────────────────────── # --------------------------------------------------------------------------- _HARD_PIPELINE: str = textwrap.dedent("""\ \"\"\"ETL pipeline — Task 3: Concurrent ingestion.\"\"\" import sqlite3 import threading # Shared connection (check_same_thread=False allows sharing) conn = sqlite3.connect("etl_target.db", check_same_thread=False) def worker(worker_id): cursor = conn.cursor() for i in range(20): # BUG: Threads share a single SQLite connection without a lock # leading to 'database is locked' errors or dropped inserts. # FIX: import threading and use a threading.Lock() around execute/commit try: cursor.execute( "INSERT INTO events (worker_id, event_no) VALUES (?, ?)", (worker_id, i) ) conn.commit() except sqlite3.OperationalError as e: print(f"Error in worker {worker_id}: {e}") threads = [] for i in range(5): t = threading.Thread(target=worker, args=(i,)) threads.append(t) t.start() for t in threads: t.join() conn.close() print("Data ingestion complete.") """) # --------------------------------------------------------------------------- # ── Environment ────────────────────────────────────────────────────────────── # --------------------------------------------------------------------------- class ETLEnv(Environment): SUPPORTS_CONCURRENT_SESSIONS: bool = True _DEFAULT_TIMEOUT_S: float = 30.0 _EDITABLE_FILES: frozenset[str] = frozenset({"pipeline.py"}) def __init__(self, difficulty: DifficultyLevel = "easy", **kwargs: Any) -> None: super().__init__(**kwargs) self.difficulty: DifficultyLevel = difficulty self._workspace: Path = Path(__file__).parent / "workspaces" / str(uuid.uuid4()) self._script_content: str = "" self._console_buffer: str = "" self._db_summary: str = "" self._done: bool = False self._grader_reward: float = 0.0 self._state: State = State(episode_id=str(uuid.uuid4()), step_count=0) def reset( self, seed: Optional[int] = None, episode_id: Optional[str] = None, difficulty: Optional[DifficultyLevel] = None, **kwargs: Any, ) -> EtlObservation: requested_difficulty = difficulty or kwargs.get("difficulty") if requested_difficulty is not None: if requested_difficulty not in ("easy", "medium", "hard"): raise ValueError("difficulty must be one of: easy, medium, hard") self.difficulty = requested_difficulty self._reset_rubric() self._state = State(episode_id=episode_id or str(uuid.uuid4()), step_count=0) self._done = False self._grader_reward = 0.0 self._console_buffer = "" if self._workspace.exists(): shutil.rmtree(self._workspace) self._workspace.mkdir(parents=True, exist_ok=True) db_path = self._workspace / "etl_target.db" pipeline_path = self._workspace / "pipeline.py" conn = sqlite3.connect(str(db_path)) # Setup varies by difficulty if self.difficulty == "easy": (self._workspace / "source_data.csv").write_text(_EASY_SOURCE_CSV, encoding="utf-8") conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL)") conn.commit() pipeline_path.write_text(_EASY_PIPELINE, encoding="utf-8") self._script_content = _EASY_PIPELINE elif self.difficulty == "medium": conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL)") conn.execute("CREATE TABLE purchases (id INTEGER PRIMARY KEY, user_id INTEGER, item TEXT NOT NULL)") # 5 users total for i, name in enumerate(["Alice", "Bob", "Charlie", "David", "Eve"], 1): conn.execute("INSERT INTO users (id, name) VALUES (?, ?)", (i, name)) # 3 purchases for users 1, 2, 3 only conn.execute("INSERT INTO purchases (user_id, item) VALUES (1, 'Book')") conn.execute("INSERT INTO purchases (user_id, item) VALUES (2, 'Pen')") conn.execute("INSERT INTO purchases (user_id, item) VALUES (3, 'Laptop')") conn.commit() pipeline_path.write_text(_MEDIUM_PIPELINE, encoding="utf-8") self._script_content = _MEDIUM_PIPELINE elif self.difficulty == "hard": conn.execute("CREATE TABLE events (id INTEGER PRIMARY KEY, worker_id INTEGER, event_no INTEGER)") conn.commit() pipeline_path.write_text(_HARD_PIPELINE, encoding="utf-8") self._script_content = _HARD_PIPELINE conn.close() self._db_summary = self._summarize_db() return EtlObservation( console_logs=self._console_buffer, current_script_content=self._script_content, db_state_summary=self._db_summary, done=self._done, reward=0.0, ) def step( self, action: EtlAction, # type: ignore[override] timeout_s: Optional[float] = None, **kwargs: Any, ) -> EtlObservation: action_type = getattr(action, "action_type", None) if action_type == "edit_code": if isinstance(action, EditCodeAction): edit_action = action else: edit_action = EditCodeAction( filepath=action.filepath, target_line=action.target_line, new_code=action.new_code, ) self._apply_edit(edit_action) elif action_type == "run_pipeline": if isinstance(action, RunPipelineAction): run_action = action else: run_action = RunPipelineAction(timeout_override_s=action.timeout_override_s) self._execute_run(run_action) elif action_type == "query_db": if isinstance(action, QueryDbAction): query_action = action else: query_action = QueryDbAction( sql_string=action.sql_string, max_rows=action.max_rows, ) self._execute_query(query_action) else: raise TypeError(f"Unrecognised action type: {type(action).__name__!r}") self._state.step_count += 1 partial_obs = EtlObservation( console_logs=self._console_buffer, current_script_content=self._script_content, db_state_summary=self._db_summary, done=self._done, reward=0.0, ) if self.rubric is not None: raw_reward = self._apply_rubric(action, partial_obs) else: raw_reward = self._grader_reward reward_obj = EtlReward(value=max(0.0, min(1.0, float(raw_reward)))) return EtlObservation( console_logs=self._console_buffer, current_script_content=self._script_content, db_state_summary=self._db_summary, done=self._done, reward=reward_obj.value, ) @property def state(self) -> State: return self._state def close(self) -> None: if self._workspace.exists(): shutil.rmtree(self._workspace, ignore_errors=True) def _apply_edit(self, action: EditCodeAction) -> None: requested_basename = Path(action.filepath).name if requested_basename not in self._EDITABLE_FILES: self._console_buffer = f"[EDIT ERROR] '{action.filepath}' is not allowed." return target_path = (self._workspace / requested_basename).resolve() if not str(target_path).startswith(str(self._workspace.resolve())): self._console_buffer = f"[EDIT ERROR] Path traversal detected." return if not target_path.exists(): self._console_buffer = f"[EDIT ERROR] File not found." return lines = target_path.read_text(encoding="utf-8").splitlines(keepends=True) if action.target_line > len(lines): self._console_buffer = f"[EDIT ERROR] target_line {action.target_line} out of range." return replacement = action.new_code if not replacement.endswith("\n"): replacement += "\n" lines[action.target_line - 1] = replacement updated_text = "".join(lines) target_path.write_text(updated_text, encoding="utf-8") self._script_content = updated_text self._console_buffer = f"[EDIT OK] Line {action.target_line} replaced." self._grader_reward = 0.0 def _execute_run(self, action: RunPipelineAction) -> None: pipeline_path = self._workspace / "pipeline.py" if not pipeline_path.exists(): self._console_buffer = "[RUN ERROR] pipeline.py missing." self._grader_reward = 0.0 return timeout = action.timeout_override_s or self._DEFAULT_TIMEOUT_S try: result = subprocess.run( [sys.executable, str(pipeline_path)], capture_output=True, text=True, timeout=timeout, cwd=str(self._workspace), ) except subprocess.TimeoutExpired: self._console_buffer = f"[RUN TIMEOUT] Exceeded {timeout:.1f}s." self._grader_reward = 0.0 return except OSError as exc: self._console_buffer = f"[RUN OS ERROR] {exc}" self._grader_reward = 0.0 return output_parts = [] if result.stdout: output_parts.append(result.stdout.rstrip()) if result.stderr: output_parts.append(f"[STDERR]\n{result.stderr.rstrip()}") exit_label = f"[Exit code: {result.returncode}]" if result.returncode != 0 else "[Exit code: 0 — OK]" output_parts.append(exit_label) self._console_buffer = "\n".join(output_parts) self._db_summary = self._summarize_db() # Score the task self._grader_reward = self._score_task() self._done = self._grader_reward >= 1.0 def _execute_query(self, action: QueryDbAction) -> None: sql = action.sql_string.strip() if not sql.upper().startswith("SELECT"): self._db_summary = "[QUERY ERROR] Only SELECT allowed." self._grader_reward = 0.0 return db_path = self._workspace / "etl_target.db" if not db_path.exists(): self._db_summary = "[QUERY ERROR] DB missing." self._grader_reward = 0.0 return try: conn = sqlite3.connect(str(db_path)) cursor = conn.execute(sql) rows = cursor.fetchmany(action.max_rows) col_names = [d[0] for d in cursor.description] if cursor.description else [] conn.close() except sqlite3.Error as exc: self._db_summary = f"[SQL ERROR] {exc}" self._grader_reward = 0.0 return if not rows: self._db_summary = f"Query OK — 0 rows.\nCols: {col_names if col_names else '(none)'}" else: header = " | ".join(col_names) separator = "-" * max(len(header), 10) data_lines = [" | ".join(str(v) for v in row) for row in rows] self._db_summary = "\n".join([header, separator] + data_lines + [f"\n({len(rows)} rows)"]) self._grader_reward = 0.0 def _summarize_db(self) -> str: db_path = self._workspace / "etl_target.db" if not db_path.exists(): return "Database missing." try: conn = sqlite3.connect(str(db_path)) summary_parts = [] if self.difficulty == "easy": (row_count,) = conn.execute("SELECT COUNT(*) FROM users").fetchone() sample = conn.execute("SELECT id, name FROM users LIMIT 3").fetchall() summary_parts.append(f"users: {row_count} rows | sample: {sample}") elif self.difficulty == "medium": (u_count,) = conn.execute("SELECT COUNT(*) FROM users").fetchone() (p_count,) = conn.execute("SELECT COUNT(*) FROM purchases").fetchone() summary_parts.append(f"users: {u_count} rows, purchases: {p_count} rows") # Check for user_summary existence safely cursor = conn.cursor() cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='user_summary'") if cursor.fetchone(): (us_count,) = conn.execute("SELECT COUNT(*) FROM user_summary").fetchone() summary_parts.append(f"user_summary: {us_count} rows") else: summary_parts.append("user_summary: not created yet") elif self.difficulty == "hard": (e_count,) = conn.execute("SELECT COUNT(*) FROM events").fetchone() summary_parts.append(f"events: {e_count} rows") conn.close() return "\n".join(summary_parts) except sqlite3.Error as exc: return f"[DB ERROR] {exc}" def _score_task(self) -> float: db_path = self._workspace / "etl_target.db" if not db_path.exists(): return 0.0 try: conn = sqlite3.connect(str(db_path)) reward = 0.0 if self.difficulty == "easy": try: rows = conn.execute("SELECT id, name FROM users ORDER BY id").fetchall() matches = sum(1 for expected in _EASY_EXPECTED_ROWS if expected in rows) reward = min(matches / len(_EASY_EXPECTED_ROWS), 1.0) except sqlite3.Error: pass elif self.difficulty == "medium": try: cursor = conn.cursor() cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='user_summary'") if cursor.fetchone(): (count,) = conn.execute("SELECT COUNT(*) FROM user_summary").fetchone() reward = min(max(count, 0) / 5.0, 1.0) except sqlite3.Error: pass elif self.difficulty == "hard": try: (count,) = conn.execute("SELECT COUNT(*) FROM events").fetchone() reward = min(count / 100.0, 1.0) except sqlite3.Error: pass conn.close() return reward except sqlite3.Error: return 0.0