@livestore/common 0.3.0-dev.22 → 0.3.0-dev.24

package/src/query-builder/astToSql.ts

@@ -0,0 +1,203 @@
+import { shouldNeverHappen } from '@livestore/utils'
+import { Schema } from '@livestore/utils/effect'
+
+import { SessionIdSymbol } from '../adapter-types.js'
+import type { DbSchema } from '../schema/mod.js'
+import type { SqlValue } from '../util.js'
+import type { QueryBuilderAst } from './api.js'
+
+// Helper functions for SQL generation
+const formatWhereClause = (
+  whereConditions: ReadonlyArray<QueryBuilderAst.Where>,
+  tableDef: DbSchema.TableDefBase,
+  bindValues: SqlValue[],
+): string => {
+  if (whereConditions.length === 0) return ''
+
+  const whereClause = whereConditions
+    .map(({ col, op, value }) => {
+      // Handle NULL values
+      if (value === null) {
+        if (op !== '=' && op !== '!=') {
+          throw new Error(`Unsupported operator for NULL value: ${op}`)
+        }
+        const opStmt = op === '=' ? 'IS' : 'IS NOT'
+        return `${col} ${opStmt} NULL`
+      }
+
+      // Get column definition and encode value
+      const colDef = tableDef.sqliteDef.columns[col]
+      if (colDef === undefined) {
+        throw new Error(`Column ${col} not found`)
+      }
+
+      // Handle array values for IN/NOT IN operators
+      const isArray = op === 'IN' || op === 'NOT IN'
+
+      if (isArray) {
+        // Verify value is an array
+        if (!Array.isArray(value)) {
+          return shouldNeverHappen(`Expected array value for ${op} operator but got`, value)
+        }
+
+        // Handle empty arrays
+        if (value.length === 0) {
+          return op === 'IN' ? '0=1' : '1=1'
+        }
+
+        const encodedValues = value.map((v) => Schema.encodeSync(colDef.schema)(v)) as SqlValue[]
+        bindValues.push(...encodedValues)
+        const placeholders = encodedValues.map(() => '?').join(', ')
+        return `${col} ${op} (${placeholders})`
+      } else {
+        const encodedValue = Schema.encodeSync(colDef.schema)(value)
+        bindValues.push(encodedValue as SqlValue)
+        return `${col} ${op} ?`
+      }
+    })
+    .join(' AND ')
+
+  return `WHERE ${whereClause}`
+}
+
+const formatReturningClause = (returning?: string[]): string => {
+  if (!returning || returning.length === 0) return ''
+  return ` RETURNING ${returning.join(', ')}`
+}
+
+export const astToSql = (ast: QueryBuilderAst): { query: string; bindValues: SqlValue[] } => {
+  const bindValues: SqlValue[] = []
+
+  // INSERT query
+  if (ast._tag === 'InsertQuery') {
+    const columns = Object.keys(ast.values)
+    const placeholders = columns.map(() => '?').join(', ')
+    const values = Object.values(Schema.encodeSync(ast.tableDef.insertSchema)(ast.values)) as SqlValue[]
+
+    bindValues.push(...values)
+
+    let query = `INSERT INTO '${ast.tableDef.sqliteDef.name}' (${columns.join(', ')}) VALUES (${placeholders})`
+
+    // Handle ON CONFLICT clause
+    if (ast.onConflict) {
+      query += ` ON CONFLICT (${ast.onConflict.target}) `
+      if (ast.onConflict.action._tag === 'ignore') {
+        query += 'DO NOTHING'
+      } else if (ast.onConflict.action._tag === 'replace') {
+        query += 'DO REPLACE'
+      } else {
+        // Handle the update record case
+        const updateValues = ast.onConflict.action.update
+        const updateCols = Object.keys(updateValues)
+        if (updateCols.length === 0) {
+          throw new Error('No update columns provided for ON CONFLICT DO UPDATE')
+        }
+
+        const updates = updateCols
+          .map((col) => {
+            const value = updateValues[col]
+            // If the value is undefined, use excluded.col
+            return value === undefined ? `${col} = excluded.${col}` : `${col} = ?`
+          })
+          .join(', ')
+
+        // Add values for the parameters
+        updateCols.forEach((col) => {
+          const value = updateValues[col]
+          if (value !== undefined) {
+            const colDef = ast.tableDef.sqliteDef.columns[col]
+            if (colDef === undefined) {
+              throw new Error(`Column ${col} not found`)
+            }
+            const encodedValue = Schema.encodeSync(colDef.schema)(value)
+            bindValues.push(encodedValue as SqlValue)
+          }
+        })
+
+        query += `DO UPDATE SET ${updates}`
+      }
+    }
+
+    query += formatReturningClause(ast.returning)
+    return { query, bindValues }
+  }
+
+  // UPDATE query
+  if (ast._tag === 'UpdateQuery') {
+    const setColumns = Object.keys(ast.values)
+    const setValues = Object.values(Schema.encodeSync(Schema.partial(ast.tableDef.schema))(ast.values))
+    bindValues.push(...setValues)
+
+    let query = `UPDATE '${ast.tableDef.sqliteDef.name}' SET ${setColumns.map((col) => `${col} = ?`).join(', ')}`
+
+    const whereClause = formatWhereClause(ast.where, ast.tableDef, bindValues)
+    if (whereClause) query += ` ${whereClause}`
+
+    query += formatReturningClause(ast.returning)
+    return { query, bindValues }
+  }
+
+  // DELETE query
+  if (ast._tag === 'DeleteQuery') {
+    let query = `DELETE FROM '${ast.tableDef.sqliteDef.name}'`
+
+    const whereClause = formatWhereClause(ast.where, ast.tableDef, bindValues)
+    if (whereClause) query += ` ${whereClause}`
+
+    query += formatReturningClause(ast.returning)
+    return { query, bindValues }
+  }
+
+  // COUNT query
+  if (ast._tag === 'CountQuery') {
+    const query = [
+      `SELECT COUNT(*) as count FROM '${ast.tableDef.sqliteDef.name}'`,
+      formatWhereClause(ast.where, ast.tableDef, bindValues),
+    ]
+      .filter((clause) => clause.length > 0)
+      .join(' ')
+
+    return { query, bindValues }
+  }
+
+  // ROW query
+  if (ast._tag === 'RowQuery') {
+    // Handle the id value by encoding it with the id column schema
+    const idColDef = ast.tableDef.sqliteDef.columns.id
+    if (idColDef === undefined) {
+      throw new Error('Column id not found for ROW query')
+    }
+
+    // NOTE we're not encoding the id if it's the session id symbol, which needs to be taken care of by the caller
+    const encodedId = ast.id === SessionIdSymbol ? ast.id : Schema.encodeSync(idColDef.schema)(ast.id)
+
+    return {
+      query: `SELECT * FROM '${ast.tableDef.sqliteDef.name}' WHERE id = ?`,
+      bindValues: [encodedId as SqlValue],
+    }
+  }
+
+  // SELECT query
+  const columnsStmt = ast.select.columns.length === 0 ? '*' : ast.select.columns.join(', ')
+  const selectStmt = `SELECT ${columnsStmt}`
+  const fromStmt = `FROM '${ast.tableDef.sqliteDef.name}'`
+  const whereStmt = formatWhereClause(ast.where, ast.tableDef, bindValues)
+
+  const orderByStmt =
+    ast.orderBy.length > 0
+      ? `ORDER BY ${ast.orderBy.map(({ col, direction }) => `${col} ${direction}`).join(', ')}`
+      : ''
+
+  const limitStmt = ast.limit._tag === 'Some' ? `LIMIT ?` : ''
+  if (ast.limit._tag === 'Some') bindValues.push(ast.limit.value)
+
+  const offsetStmt = ast.offset._tag === 'Some' ? `OFFSET ?` : ''
+  if (ast.offset._tag === 'Some') bindValues.push(ast.offset.value)
+
+  const query = [selectStmt, fromStmt, whereStmt, orderByStmt, offsetStmt, limitStmt]
+    .map((clause) => clause.trim())
+    .filter((clause) => clause.length > 0)
+    .join(' ')
+
+  return { query, bindValues }
+}

package/src/query-builder/impl.test.ts

@@ -226,6 +226,110 @@ describe('query builder', () => {
       `)
     })
   })
+
+  describe('write operations', () => {
+    it('should handle INSERT queries', () => {
+      expect(db.todos.insert({ id: '123', text: 'Buy milk', status: 'active' }).asSql()).toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "123",
+            "Buy milk",
+            "active",
+          ],
+          "query": "INSERT INTO 'todos' (id, text, status) VALUES (?, ?, ?)",
+        }
+      `)
+    })
+
+    it('should handle UPDATE queries', () => {
+      expect(db.todos.update({ status: 'completed' }).where({ id: '123' }).asSql()).toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "completed",
+            "123",
+          ],
+          "query": "UPDATE 'todos' SET status = ? WHERE id = ?",
+        }
+      `)
+    })
+
+    it('should handle DELETE queries', () => {
+      expect(db.todos.delete().where({ status: 'completed' }).asSql()).toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "completed",
+          ],
+          "query": "DELETE FROM 'todos' WHERE status = ?",
+        }
+      `)
+    })
+
+    it('should handle INSERT with ON CONFLICT', () => {
+      expect(db.todos.insert({ id: '123', text: 'Buy milk', status: 'active' }).onConflict('id', 'ignore').asSql())
+        .toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "123",
+            "Buy milk",
+            "active",
+          ],
+          "query": "INSERT INTO 'todos' (id, text, status) VALUES (?, ?, ?) ON CONFLICT (id) DO NOTHING",
+        }
+      `)
+
+      expect(
+        db.todos
+          .insert({ id: '123', text: 'Buy milk', status: 'active' })
+          .onConflict('id', 'update', { text: 'Buy soy milk', status: 'active' })
+          .asSql(),
+      ).toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "123",
+            "Buy milk",
+            "active",
+            "Buy soy milk",
+            "active",
+          ],
+          "query": "INSERT INTO 'todos' (id, text, status) VALUES (?, ?, ?) ON CONFLICT (id) DO UPDATE SET text = ?, status = ?",
+        }
+      `)
+    })
+
+    it('should handle RETURNING clause', () => {
+      expect(db.todos.insert({ id: '123', text: 'Buy milk', status: 'active' }).returning('id').asSql())
+        .toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "123",
+            "Buy milk",
+            "active",
+          ],
+          "query": "INSERT INTO 'todos' (id, text, status) VALUES (?, ?, ?) RETURNING id",
+        }
+      `)
+
+      expect(db.todos.update({ status: 'completed' }).where({ id: '123' }).returning('id').asSql())
+        .toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "completed",
+            "123",
+          ],
+          "query": "UPDATE 'todos' SET status = ? WHERE id = ? RETURNING id",
+        }
+      `)
+
+      expect(db.todos.delete().where({ status: 'completed' }).returning('id').asSql()).toMatchInlineSnapshot(`
+        {
+          "bindValues": [
+            "completed",
+          ],
+          "query": "DELETE FROM 'todos' WHERE status = ? RETURNING id",
+        }
+      `)
+    })
+  })
 })
 
 // TODO nested queries

package/src/query-builder/impl.ts

@@ -1,9 +1,11 @@
-import { Option, Predicate, Schema } from '@livestore/utils/effect'
+import { casesHandled } from '@livestore/utils'
+import { Match, Option, Predicate, Schema } from '@livestore/utils/effect'
 
 import type { QueryInfo } from '../query-info.js'
 import type { DbSchema } from '../schema/mod.js'
 import type { QueryBuilder, QueryBuilderAst } from './api.js'
 import { QueryBuilderAstSymbol, TypeId } from './api.js'
+import { astToSql } from './astToSql.js'
 
 export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBase>(
   tableDef: TTableDef,
@@ -12,7 +14,7 @@ export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBas
   const api = {
     // eslint-disable-next-line prefer-arrow/prefer-arrow-functions
     select() {
-      assertQueryBuilderAst(ast)
+      assertSelectQueryBuilderAst(ast)
 
       // eslint-disable-next-line prefer-rest-params
       const params = [...arguments]
@@ -36,8 +38,9 @@ export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBas
       }) as any
     },
     // eslint-disable-next-line prefer-arrow/prefer-arrow-functions
-    where() {
-      if (isRowQuery(ast)) return invalidQueryBuilder()
+    where: function () {
+      if (ast._tag === 'InsertQuery') return invalidQueryBuilder('Cannot use where with insert')
+      if (ast._tag === 'RowQuery') return invalidQueryBuilder('Cannot use where with row')
 
       if (arguments.length === 1) {
         // eslint-disable-next-line prefer-rest-params
@@ -50,24 +53,45 @@ export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBas
               : { col, op: '=', value },
           )
 
-        return makeQueryBuilder(tableDef, {
-          ...ast,
-          where: [...ast.where, ...newOps],
-        }) as any
+        switch (ast._tag) {
+          case 'CountQuery':
+          case 'SelectQuery':
+          case 'UpdateQuery':
+          case 'DeleteQuery': {
+            return makeQueryBuilder(tableDef, {
+              ...ast,
+              where: [...ast.where, ...newOps],
+            }) as any
+          }
+          default: {
+            return casesHandled(ast)
+          }
+        }
       }
 
       // eslint-disable-next-line prefer-rest-params
       const [col, opOrValue, valueOrUndefined] = arguments
       const op = valueOrUndefined === undefined ? '=' : opOrValue
       const value = valueOrUndefined === undefined ? opOrValue : valueOrUndefined
-      return makeQueryBuilder(tableDef, {
-        ...ast,
-        where: [...ast.where, { col, op, value }],
-      })
+
+      switch (ast._tag) {
+        case 'CountQuery':
+        case 'SelectQuery':
+        case 'UpdateQuery':
+        case 'DeleteQuery': {
+          return makeQueryBuilder(tableDef, {
+            ...ast,
+            where: [...ast.where, { col, op, value }],
+          }) as any
+        }
+        default: {
+          return casesHandled(ast)
+        }
+      }
     },
     // eslint-disable-next-line prefer-arrow/prefer-arrow-functions
     orderBy() {
-      assertQueryBuilderAst(ast)
+      assertSelectQueryBuilderAst(ast)
 
       if (arguments.length === 0 || arguments.length > 2) return invalidQueryBuilder()
 
@@ -89,12 +113,12 @@ export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBas
       }) as any
     },
     limit: (limit) => {
-      assertQueryBuilderAst(ast)
+      assertSelectQueryBuilderAst(ast)
 
       return makeQueryBuilder(tableDef, { ...ast, limit: Option.some(limit) })
     },
     offset: (offset) => {
-      assertQueryBuilderAst(ast)
+      assertSelectQueryBuilderAst(ast)
 
       return makeQueryBuilder(tableDef, { ...ast, offset: Option.some(offset) })
     },
@@ -102,17 +126,18 @@ export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBas
       if (isRowQuery(ast)) return invalidQueryBuilder()
 
       return makeQueryBuilder(tableDef, {
-        ...ast,
+        _tag: 'CountQuery',
+        tableDef,
+        where: [],
         resultSchema: Schema.Struct({ count: Schema.Number }).pipe(
           Schema.pluck('count'),
           Schema.Array,
           Schema.headOrElse(),
         ),
-        _tag: 'CountQuery',
       })
     },
     first: (options) => {
-      assertQueryBuilderAst(ast)
+      assertSelectQueryBuilderAst(ast)
 
       if (ast.limit._tag === 'Some') return invalidQueryBuilder(`.first() can't be called after .limit()`)
 
@@ -149,6 +174,65 @@ export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBas
         insertValues,
       }) as any
     },
+    insert: (values) => {
+      return makeQueryBuilder(tableDef, {
+        _tag: 'InsertQuery',
+        tableDef,
+        values: values as any,
+        onConflict: undefined,
+        returning: undefined,
+        resultSchema: Schema.Void,
+      }) as any
+    },
+    onConflict: (target: string, action: 'ignore' | 'replace' | 'update', updateValues?: Record<string, unknown>) => {
+      assertInsertQueryBuilderAst(ast)
+
+      const onConflict = Match.value(action).pipe(
+        Match.when('ignore', () => ({ target, action: { _tag: 'ignore' } }) satisfies QueryBuilderAst.OnConflict),
+        Match.when('replace', () => ({ target, action: { _tag: 'replace' } }) satisfies QueryBuilderAst.OnConflict),
+        Match.when(
+          'update',
+          () => ({ target, action: { _tag: 'update', update: updateValues! } }) satisfies QueryBuilderAst.OnConflict,
+        ),
+        Match.exhaustive,
+      )
+
+      return makeQueryBuilder(tableDef, {
+        ...ast,
+        onConflict,
+      }) as any
+    },
+
+    returning: (...columns) => {
+      assertWriteQueryBuilderAst(ast)
+
+      return makeQueryBuilder(tableDef, {
+        ...ast,
+        returning: columns,
+        resultSchema: tableDef.schema.pipe(Schema.pick(...columns), Schema.Array),
+      }) as any
+    },
+
+    update: (values) => {
+      return makeQueryBuilder(tableDef, {
+        _tag: 'UpdateQuery',
+        tableDef,
+        values: values as any,
+        where: [],
+        returning: undefined,
+        resultSchema: Schema.Void,
+      }) as any
+    },
+
+    delete: () => {
+      return makeQueryBuilder(tableDef, {
+        _tag: 'DeleteQuery',
+        tableDef,
+        where: [],
+        returning: undefined,
+        resultSchema: Schema.Void,
+      }) as any
+    },
   } satisfies QueryBuilder.ApiFull<TResult, TTableDef, never, QueryInfo.None>
 
   return {
@@ -167,94 +251,38 @@ export const makeQueryBuilder = <TResult, TTableDef extends DbSchema.TableDefBas
   } satisfies QueryBuilder<TResult, TTableDef>
 }
 
-const emptyAst = (tableDef: DbSchema.TableDefBase) =>
-  ({
-    _tag: 'SelectQuery',
-    columns: [],
-    pickFirst: false,
-    select: { columns: [] },
-    orderBy: [],
-    offset: Option.none(),
-    limit: Option.none(),
-    tableDef,
-    where: [],
-    resultSchemaSingle: tableDef.schema,
-  }) satisfies QueryBuilderAst
-
-const astToSql = (ast: QueryBuilderAst) => {
-  if (isRowQuery(ast)) {
-    // TODO
-    return { query: `SELECT * FROM '${ast.tableDef.sqliteDef.name}' WHERE id = ?`, bindValues: [ast.id as TODO] }
+const emptyAst = (tableDef: DbSchema.TableDefBase): QueryBuilderAst.SelectQuery => ({
+  _tag: 'SelectQuery',
+  columns: [],
+  pickFirst: false,
+  select: { columns: [] },
+  orderBy: [],
+  offset: Option.none(),
+  limit: Option.none(),
+  tableDef,
+  where: [],
+  resultSchemaSingle: tableDef.schema,
+})
+
+// Helper functions
+// eslint-disable-next-line prefer-arrow/prefer-arrow-functions
+function assertSelectQueryBuilderAst(ast: QueryBuilderAst): asserts ast is QueryBuilderAst.SelectQuery {
+  if (ast._tag !== 'SelectQuery') {
+    throw new Error('Expected SelectQuery but got ' + ast._tag)
   }
+}
 
-  const bindValues: unknown[] = []
-
-  // TODO bind values
-  const whereStmt =
-    ast.where.length > 0
-      ? `WHERE ${ast.where
-          .map(({ col, op, value }) => {
-            if (value === null) {
-              if (op !== '=' && op !== '!=') {
-                throw new Error(`Unsupported operator for NULL value: ${op}`)
-              }
-              const opStmt = op === '=' ? 'IS' : 'IS NOT'
-              return `${col} ${opStmt} NULL`
-            } else {
-              const colDef = ast.tableDef.sqliteDef.columns[col]
-              if (colDef === undefined) {
-                throw new Error(`Column ${col} not found`)
-              }
-              const isArray = op === 'IN' || op === 'NOT IN'
-              const colSchema = isArray ? Schema.Array(colDef.schema) : colDef.schema
-              const encodedValue = Schema.encodeSync(colSchema)(value)
-
-              if (isArray) {
-                bindValues.push(...encodedValue)
-                const placeholders = Array.from({ length: encodedValue.length }, () => '?').join(', ')
-                return `${col} ${op} (${placeholders})`
-              } else {
-                bindValues.push(encodedValue)
-                return `${col} ${op} ?`
-              }
-            }
-          })
-          .join(' AND ')}`
-      : ''
-
-  if (ast._tag === 'CountQuery') {
-    const selectFromStmt = `SELECT COUNT(*) as count FROM '${ast.tableDef.sqliteDef.name}'`
-    const query = [selectFromStmt, whereStmt].filter((_) => _.length > 0).join(' ')
-    return { query, bindValues }
+// eslint-disable-next-line prefer-arrow/prefer-arrow-functions
+function assertInsertQueryBuilderAst(ast: QueryBuilderAst): asserts ast is QueryBuilderAst.InsertQuery {
+  if (ast._tag !== 'InsertQuery') {
+    throw new Error('Expected InsertQuery but got ' + ast._tag)
   }
-  const columnsStmt = ast.select.columns.length === 0 ? '*' : ast.select.columns.join(', ')
-  const selectStmt = `SELECT ${columnsStmt}`
-  const fromStmt = `FROM '${ast.tableDef.sqliteDef.name}'`
-
-  const orderByStmt =
-    ast.orderBy.length > 0
-      ? `ORDER BY ${ast.orderBy.map(({ col, direction }) => `${col} ${direction}`).join(', ')}`
-      : ''
-
-  const limitStmt = ast.limit._tag === 'Some' ? `LIMIT ?` : ''
-  if (ast.limit._tag === 'Some') bindValues.push(ast.limit.value)
-
-  const offsetStmt = ast.offset._tag === 'Some' ? `OFFSET ?` : ''
-  if (ast.offset._tag === 'Some') bindValues.push(ast.offset.value)
-
-  const query = [selectStmt, fromStmt, whereStmt, orderByStmt, offsetStmt, limitStmt]
-    .map((_) => _.trim())
-    .filter((_) => _.length > 0)
-    .join(' ')
-
-  // TODO bind values
-  return { query, bindValues }
 }
 
 // eslint-disable-next-line prefer-arrow/prefer-arrow-functions
-function assertQueryBuilderAst(ast: QueryBuilderAst): asserts ast is QueryBuilderAst.SelectQuery {
-  if (ast._tag !== 'SelectQuery') {
-    throw new Error('Expected SelectQuery but got ' + ast._tag)
+function assertWriteQueryBuilderAst(ast: QueryBuilderAst): asserts ast is QueryBuilderAst.WriteQuery {
+  if (ast._tag !== 'InsertQuery' && ast._tag !== 'UpdateQuery' && ast._tag !== 'DeleteQuery') {
+    throw new Error('Expected WriteQuery but got ' + ast._tag)
   }
 }
 
@@ -264,22 +292,38 @@ export const invalidQueryBuilder = (msg?: string) => {
   throw new Error('Invalid query builder' + (msg ? `: ${msg}` : ''))
 }
 
-export const getResultSchema = (qb: QueryBuilder<any, any, any>) => {
+export const getResultSchema = (qb: QueryBuilder<any, any, any>): Schema.Schema<any> => {
   const queryAst = qb[QueryBuilderAstSymbol]
-  if (queryAst._tag === 'SelectQuery') {
-    const arraySchema = Schema.Array(queryAst.resultSchemaSingle)
-    if (queryAst.pickFirst !== false) {
-      return arraySchema.pipe(Schema.headOrElse(queryAst.pickFirst.fallback))
+  switch (queryAst._tag) {
+    case 'SelectQuery': {
+      const arraySchema = Schema.Array(queryAst.resultSchemaSingle)
+      if (queryAst.pickFirst !== false) {
+        return arraySchema.pipe(Schema.headOrElse(queryAst.pickFirst.fallback))
+      }
+
+      return arraySchema
     }
+    case 'CountQuery': {
+      return Schema.Struct({ count: Schema.Number }).pipe(Schema.pluck('count'), Schema.Array, Schema.headOrElse())
+    }
+    case 'InsertQuery':
+    case 'UpdateQuery':
+    case 'DeleteQuery': {
+      // For write operations with RETURNING clause, we need to return the appropriate schema
+      if (queryAst.returning && queryAst.returning.length > 0) {
+        // Create a schema for the returned columns
+        return queryAst.tableDef.schema.pipe(Schema.pick(...queryAst.returning), Schema.Array)
+      }
 
-    return arraySchema
-  } else if (queryAst._tag === 'CountQuery') {
-    return Schema.Struct({ count: Schema.Number }).pipe(Schema.pluck('count'), Schema.Array, Schema.headOrElse())
-  } else {
-    if (queryAst.tableDef.options.isSingleColumn) {
-      return queryAst.tableDef.schema.pipe(Schema.pluck('value'), Schema.Array, Schema.headOrElse())
-    } else {
-      return queryAst.tableDef.schema.pipe(Schema.Array, Schema.headOrElse())
+      // For write operations without RETURNING, the result is the number of affected rows
+      return Schema.Number
+    }
+    default: {
+      if (queryAst.tableDef.options.isSingleColumn) {
+        return queryAst.tableDef.schema.pipe(Schema.pluck('value'), Schema.Array, Schema.headOrElse())
+      } else {
+        return queryAst.tableDef.schema.pipe(Schema.Array, Schema.headOrElse())
+      }
     }
   }
 }

package/src/query-builder/mod.ts

@@ -7,4 +7,11 @@ export * from './impl.js'
  * - Close abstraction to SQLite to provide a simple & convenient API with predictable behaviour
  * - Use table schema definitions to parse, map & validate query results
  * - Implementation detail: Separate type-level & AST-based runtime implementation
+ *
+ * Currently not supported (not exhaustive list):
+ * - Assumes a `id` column as primary key
+ * - Composite primary keys
+ *
+ * Other known limitations
+ * - Doesn't exclude all invalid query patterns on type level `e.g. `db.todos.returning('id')`
  */