feat(sql) Add SQL support for UPDATE, INSERT, UPSERT

protocolstardust · protocolstardust · commit 0df1b3d47339 · 2026-01-17T00:52:18.000+01:00
diff --git a/docs/docs/content/documentation/plugins/sql.md b/docs/docs/content/documentation/plugins/sql.md
@@ -147,11 +147,124 @@ employees.sql("SELECT name, salary / 12 AS monthly_salary FROM self")
 employees.sql("SELECT name, salary + 5000 AS adjusted_salary FROM self")
 ```
 
+## UPDATE Statements
+
+Modify existing rows using UPDATE with SET and optional WHERE:
+
+```python
+# Update all rows
+employees.sql("UPDATE self SET salary = 50000.0")
+
+# Update with WHERE clause
+employees.sql("UPDATE self SET salary = 100000.0 WHERE level = 'senior'")
+
+# Update multiple columns
+employees.sql("UPDATE self SET salary = 75000.0, level = 'mid' WHERE id = 5")
+
+# Update with expressions (using column values)
+employees.sql("UPDATE self SET salary = salary * 1.1 WHERE rating > 4.0")
+
+# Update with complex WHERE
+employees.sql("UPDATE self SET salary = salary + 5000.0 WHERE (dept = 'eng' OR dept = 'sales') AND years > 3")
+```
+
+**Note:** UPDATE returns the modified table. The original table is not mutated.
+
+## INSERT Statements
+
+Add new rows to a table using INSERT with VALUES:
+
+```python
+>>> employees = Table({
+...     "id": Vector([1, 2], ray_type=I64),
+...     "name": Vector(["Alice", "Bob"], ray_type=Symbol),
+...     "salary": Vector([50000.0, 60000.0], ray_type=F64),
+... })
+
+# Insert single row with column names
+>>> result = employees.sql("INSERT INTO self (id, name, salary) VALUES (3, 'Charlie', 70000.0)")
+>>> print(len(result))
+3
+
+# Insert multiple rows
+>>> result = employees.sql("""
+...     INSERT INTO self (id, name, salary)
+...     VALUES (4, 'Diana', 55000.0), (5, 'Eve', 65000.0)
+... """)
+>>> print(len(result))
+5
+
+# Insert without column names (values must match table column order)
+>>> result = employees.sql("INSERT INTO self VALUES (6, 'Frank', 72000.0)")
+```
+
+**Note:** INSERT returns a new table with the added rows. The original table is not mutated.
+
+## UPSERT (INSERT ... ON CONFLICT)
+
+Perform upsert operations (insert or update) using the `ON CONFLICT` clause:
+
+```python
+>>> products = Table({
+...     "id": Vector([1, 2], ray_type=I64),
+...     "name": Vector(["Widget", "Gadget"], ray_type=Symbol),
+...     "price": Vector([10.0, 20.0], ray_type=F64),
+... })
+
+# Update existing row (id=1 exists)
+>>> result = products.sql("""
+...     INSERT INTO self (id, name, price)
+...     VALUES (1, 'Widget Pro', 15.0)
+...     ON CONFLICT (id) DO UPDATE
+... """)
+>>> print(result["name"][0].value)
+Widget Pro
+
+# Insert new row (id=3 doesn't exist)
+>>> result = products.sql("""
+...     INSERT INTO self (id, name, price)
+...     VALUES (3, 'Gizmo', 30.0)
+...     ON CONFLICT (id) DO UPDATE
+... """)
+>>> print(len(result))
+3
+
+# Upsert multiple rows at once
+>>> result = products.sql("""
+...     INSERT INTO self (id, name, price)
+...     VALUES (1, 'Widget Updated', 12.0), (4, 'Doohickey', 25.0)
+...     ON CONFLICT (id) DO UPDATE
+... """)
+```
+
+### Composite Keys
+
+Use multiple columns as the conflict key by listing them in the `ON CONFLICT` clause:
+
+```python
+>>> inventory = Table({
+...     "region": Vector(["US", "EU"], ray_type=Symbol),
+...     "sku": Vector(["A001", "A001"], ray_type=Symbol),
+...     "quantity": Vector([100, 50], ray_type=I64),
+... })
+
+# Use (region, sku) as composite key
+>>> result = inventory.sql("""
+...     INSERT INTO self (region, sku, quantity)
+...     VALUES ('US', 'A001', 150)
+...     ON CONFLICT (region, sku) DO UPDATE
+... """)
+```
+
+**Important:** The conflict key columns must match the first N columns in the INSERT column list (in order). This is required because Rayforce uses positional key columns.
+
+**Note:** UPSERT returns a new table with the changes applied. The original table is not mutated. `ON CONFLICT DO NOTHING` is not supported.
+
 ## Limitations
 
 The current SQL implementation supports common query patterns but has some limitations:
 
-- Only `SELECT` statements are supported (no `INSERT`, `UPDATE`, `DELETE`)
+- `DELETE` statements are not supported
 - `JOIN` operations are not yet supported via SQL (use the native `.inner_join()`, `.left_join()` methods)
 - Subqueries are not supported
 - `HAVING` clause is not supported
diff --git a/rayforce/plugins/sql.py b/rayforce/plugins/sql.py
@@ -77,15 +77,45 @@ class ParsedSelect:
     order_by: list[tuple[str, bool]] = field(default_factory=list)  # (col, is_desc)
 
 
+@dataclass
+class ParsedUpdate:
+    assignments: dict[str, ParsedExpr]
+    where_clause: ParsedExpr | None = None
+
+
+@dataclass
+class ParsedInsert:
+    columns: list[str] | None
+    values: list[list[ParsedExpr]]
+
+
+@dataclass
+class ParsedUpsert:
+    columns: list[str] | None
+    values: list[list[ParsedExpr]]
+    key_columns: int
+
+
+ParsedQuery = ParsedSelect | ParsedUpdate | ParsedInsert | ParsedUpsert
+
+
 class SQLParser:
-    def parse(self, sql: str) -> ParsedSelect:
+    def parse(self, sql: str) -> ParsedQuery:
         sqlglot = _ensure_sqlglot()
         ast = sqlglot.parse_one(sql)
 
-        if ast.key != "select":
-            raise ValueError(f"Only SELECT statements are supported, got: {ast.key}")
-
-        return self._parse_select(ast)
+        if ast.key == "select":
+            return self._parse_select(ast)
+        if ast.key == "update":
+            return self._parse_update(ast)
+        if ast.key == "insert":
+            if ast.args.get("conflict"):  # on conflict
+                return self._parse_upsert(ast)
+            return self._parse_insert(ast)
+
+        raise ValueError(
+            f"Only SELECT, UPDATE, INSERT, and UPSERT statements are supported, got: {ast.key}"
+        )
 
     def _parse_select(self, node: exp.Select) -> ParsedSelect:
         import sqlglot.expressions as exp
@@ -127,6 +157,97 @@ def _parse_select(self, node: exp.Select) -> ParsedSelect:
             order_by=order_by,
         )
 
+    def _parse_update(self, node: exp.Expression) -> ParsedUpdate:
+        import sqlglot.expressions as exp
+
+        assignments: dict[str, ParsedExpr] = {}
+        for expr in node.expressions:
+            if isinstance(expr, exp.EQ):
+                col_name = expr.this.name if isinstance(expr.this, exp.Column) else str(expr.this)
+                assignments[col_name] = self._parse_expr(expr.expression)
+
+        where_clause = None
+        if node.args.get("where"):
+            where_clause = self._parse_expr(node.args["where"].this)
+
+        return ParsedUpdate(assignments=assignments, where_clause=where_clause)
+
+    def _parse_insert(self, node: exp.Expression) -> ParsedInsert:
+        import sqlglot.expressions as exp
+
+        columns: list[str] | None = None
+        if hasattr(node.this, "expressions") and node.this.expressions:
+            columns = [
+                col.name if isinstance(col, exp.Column) else str(col)
+                for col in node.this.expressions
+            ]
+
+        values: list[list[ParsedExpr]] = []
+        values_clause = node.args.get("expression")
+        if values_clause and hasattr(values_clause, "expressions"):
+            for row_tuple in values_clause.expressions:
+                if hasattr(row_tuple, "expressions"):
+                    row = [self._parse_expr(val) for val in row_tuple.expressions]
+                    values.append(row)
+
+        if not values:
+            raise ValueError("INSERT statement must have VALUES")
+
+        return ParsedInsert(columns=columns, values=values)
+
+    def _parse_upsert(self, node: exp.Expression) -> ParsedUpsert:
+        import sqlglot.expressions as exp
+
+        columns: list[str] | None = None
+        if hasattr(node.this, "expressions") and node.this.expressions:
+            columns = [
+                col.name if isinstance(col, exp.Column) else str(col)
+                for col in node.this.expressions
+            ]
+
+        values: list[list[ParsedExpr]] = []
+        values_clause = node.args.get("expression")
+        if values_clause and hasattr(values_clause, "expressions"):
+            for row_tuple in values_clause.expressions:
+                if hasattr(row_tuple, "expressions"):
+                    row = [self._parse_expr(val) for val in row_tuple.expressions]
+                    values.append(row)
+
+        if not values:
+            raise ValueError("UPSERT statement must have VALUES")
+
+        conflict = node.args.get("conflict")
+        conflict_keys: list[str] = []
+        if conflict:
+            action = conflict.args.get("action")
+            if action and str(action) == "DO NOTHING":
+                raise ValueError("ON CONFLICT DO NOTHING is not supported, use DO UPDATE")
+
+            # Get conflict key columns
+            keys = conflict.args.get("conflict_keys", [])
+            conflict_keys = [k.name if hasattr(k, "name") else str(k) for k in keys]
+
+            # MySQL-style ON DUPLICATE KEY doesn't have conflict_keys
+            if conflict.args.get("duplicate") and not conflict_keys:
+                raise ValueError(
+                    "ON DUPLICATE KEY UPDATE requires explicit key columns. "
+                    "Use: ON CONFLICT (col1, col2) DO UPDATE"
+                )
+
+        if not conflict_keys:
+            raise ValueError("UPSERT requires ON CONFLICT (key_columns) clause")
+
+        # Validate that conflict keys match the first N columns
+        if columns:
+            for i, key in enumerate(conflict_keys):
+                if i >= len(columns) or columns[i] != key:
+                    raise ValueError(
+                        f"Conflict key '{key}' must match the first {len(conflict_keys)} columns. "
+                        f"Expected '{columns[i] if i < len(columns) else 'N/A'}' at position {i}"
+                    )
+
+        return ParsedUpsert(columns=columns, values=values, key_columns=len(conflict_keys))
+
     def _parse_expr(self, node: exp.Expression) -> ParsedExpr:
         import sqlglot.expressions as exp
 
@@ -240,7 +361,18 @@ def _parse_expr(self, node: exp.Expression) -> ParsedExpr:
 
 
 class SQLCompiler:
-    def compile(self, parsed: ParsedSelect, table: Table) -> Table:
+    def compile(self, parsed: ParsedQuery, table: Table) -> Table:
+        if isinstance(parsed, ParsedSelect):
+            return self._compile_select(parsed, table)
+        if isinstance(parsed, ParsedUpdate):
+            return self._compile_update(parsed, table)
+        if isinstance(parsed, ParsedInsert):
+            return self._compile_insert(parsed, table)
+        if isinstance(parsed, ParsedUpsert):
+            return self._compile_upsert(parsed, table)
+        raise ValueError(f"Unsupported query type: {type(parsed).__name__}")
+
+    def _compile_select(self, parsed: ParsedSelect, table: Table) -> Table:
         select_args: list[str] = []
         select_kwargs: dict[str, t.Any] = {}
 
@@ -274,6 +406,61 @@ def compile(self, parsed: ParsedSelect, table: Table) -> Table:
 
         return query.execute()
 
+    def _compile_update(self, parsed: ParsedUpdate, table: Table) -> Table:
+        update_kwargs: dict[str, t.Any] = {}
+        for col_name, expr in parsed.assignments.items():
+            compiled = self._compile_expr(expr)
+            update_kwargs[col_name] = compiled
+
+        query = table.update(**update_kwargs)
+
+        if parsed.where_clause:
+            where_expr = self._compile_expr(parsed.where_clause)
+            query = query.where(where_expr)
+
+        return query.execute()
+
+    def _compile_insert(self, parsed: ParsedInsert, table: Table) -> Table:
+        compiled_rows: list[list[t.Any]] = []
+        for row in parsed.values:
+            compiled_row = [self._compile_expr(val) for val in row]
+            compiled_rows.append(compiled_row)
+
+        if parsed.columns:
+            # INSERT with column names: use kwargs style
+            insert_kwargs: dict[str, list[t.Any]] = {col: [] for col in parsed.columns}
+            for row in compiled_rows:
+                for i, col in enumerate(parsed.columns):
+                    insert_kwargs[col].append(row[i])
+            return table.insert(**insert_kwargs).execute()
+
+        if len(compiled_rows) == 1:
+            return table.insert(*compiled_rows[0]).execute()
+
+        num_cols = len(compiled_rows[0])  # transpose
+        col_values = [[row[i] for row in compiled_rows] for i in range(num_cols)]
+        return table.insert(*col_values).execute()
+
+    def _compile_upsert(self, parsed: ParsedUpsert, table: Table) -> Table:
+        compiled_rows: list[list[t.Any]] = []
+        for row in parsed.values:
+            compiled_row = [self._compile_expr(val) for val in row]
+            compiled_rows.append(compiled_row)
+
+        if parsed.columns:
+            upsert_kwargs: dict[str, list[t.Any]] = {col: [] for col in parsed.columns}
+            for row in compiled_rows:
+                for i, col in enumerate(parsed.columns):
+                    upsert_kwargs[col].append(row[i])
+            return table.upsert(**upsert_kwargs, key_columns=parsed.key_columns).execute()
+
+        if len(compiled_rows) == 1:
+            return table.upsert(*compiled_rows[0], key_columns=parsed.key_columns).execute()
+
+        num_cols = len(compiled_rows[0])  # transpose
+        col_values = [[row[i] for row in compiled_rows] for i in range(num_cols)]
+        return table.upsert(*col_values, key_columns=parsed.key_columns).execute()
+
     def _compile_expr(self, expr: ParsedExpr) -> t.Any:
         from rayforce.types.table import Column
 
diff --git a/tests/plugins/test_sql.py b/tests/plugins/test_sql.py
