on-zero 0.0.1

package/src/createZeroServer.ts

@@ -0,0 +1,226 @@
+import { mustGetQuery } from '@rocicorp/zero'
+import { PushProcessor } from '@rocicorp/zero/pg'
+import { handleQueryRequest as zeroHandleQueryRequest } from '@rocicorp/zero/server'
+import { zeroNodePg } from '@rocicorp/zero/server/adapters/pg'
+import { assertString, randomId } from '@take-out/helpers'
+import { Pool } from 'pg'
+
+import { createPermissions } from './createPermissions'
+import { createMutators } from './helpers/createMutators'
+import { isInZeroMutation, mutatorContext } from './helpers/mutatorContext'
+import { setCustomQueries } from './run'
+import { setAuthData, setSchema } from './state'
+import { setRunner } from './zeroRunner'
+
+import type {
+  AsyncAction,
+  AuthData,
+  GenericModels,
+  GetZeroMutators,
+  Transaction,
+} from './types'
+import type {
+  AnyQueryRegistry,
+  HumanReadable,
+  Query,
+  Schema as ZeroSchema,
+} from '@rocicorp/zero'
+import type { TransactionProviderInput } from '@rocicorp/zero/pg'
+
+export function createZeroServer<
+  Schema extends ZeroSchema,
+  Models extends GenericModels,
+  ServerActions extends Record<string, unknown>,
+>({
+  createServerActions,
+  database,
+  schema,
+  models,
+  queries,
+}: {
+  /**
+   * The DB connection string, same as ZERO_UPSTREAM_DB
+   */
+  database: string
+  schema: Schema
+  models: Models
+  createServerActions: () => ServerActions
+  queries?: AnyQueryRegistry
+}) {
+  setSchema(schema)
+
+  const dbString = assertString(database, `createZeroServer "database"`)
+
+  const zeroDb = zeroNodePg(
+    schema,
+    new Pool({
+      connectionString: dbString,
+      // handle self-signed certificates in production
+      ssl: dbString.includes('sslmode=require')
+        ? { rejectUnauthorized: false }
+        : undefined,
+    })
+  )
+
+  const permissions = createPermissions<Schema>({
+    environment: 'server',
+    schema,
+  })
+
+  const processor = new PushProcessor(zeroDb)
+
+  const handleMutationRequest = async ({
+    authData,
+    request,
+    skipAsyncTasks,
+  }: {
+    authData: AuthData | null
+    request: Request
+    skipAsyncTasks?: boolean
+  }) => {
+    // since mutations do DB work in transaction, avoid any async tasks during
+    const asyncTasks: AsyncAction[] = []
+
+    const mutators = createMutators({
+      asyncTasks,
+      can: permissions.can,
+      createServerActions,
+      environment: 'server',
+      models,
+      authData,
+    })
+
+    // @ts-expect-error type is ok but config in monorepo
+    const response = await processor.process(mutators, request)
+
+    // now finish
+    if (!skipAsyncTasks && asyncTasks.length) {
+      const id = randomId()
+      console.info(`[push] complete, running async tasks ${asyncTasks.length} id ${id}`)
+      Promise.all(asyncTasks.map((task) => task()))
+        .then(() => {
+          console.info(`[push] async tasks complete ${id}`)
+        })
+        .catch((err) => {
+          console.error(`[push] error: async tasks failed 😞`, err)
+        })
+    }
+
+    return {
+      response,
+      asyncTasks,
+    }
+  }
+
+  const handleQueryRequest = async ({
+    authData,
+    request,
+  }: {
+    authData: AuthData | null
+    request: Request
+  }) => {
+    if (!queries) {
+      throw new Error(
+        'No queries registered with createZeroServer. ' +
+          'Pass the syncedQueries registry to createZeroServer via the queries option.'
+      )
+    }
+
+    // set authData globally for permission checks in query functions
+    setAuthData(authData || ({} as AuthData))
+
+    const response = await zeroHandleQueryRequest(
+      (name, args) => {
+        const query = (mustGetQuery as any)(queries, name)
+        return query.fn({ args, ctx: authData })
+      },
+      schema,
+      request
+    )
+
+    return {
+      response,
+    }
+  }
+
+  const mutate = async (
+    run: (tx: Transaction, mutators: GetZeroMutators<Models>) => Promise<void>,
+    authData?: Pick<AuthData, 'email' | 'id'> & Partial<AuthData>
+  ) => {
+    const asyncTasks: Array<() => Promise<void>> = []
+
+    const mutators = createMutators({
+      models,
+      environment: 'server',
+      asyncTasks,
+      authData: {
+        id: '',
+        email: 'admin@start.chat',
+        role: 'admin',
+        ...authData,
+      },
+      createServerActions,
+      can: permissions.can,
+    })
+
+    await transaction(async (tx) => {
+      await run(tx, mutators)
+    })
+
+    await Promise.all(asyncTasks.map((t) => t()))
+  }
+
+  async function transaction<
+    CB extends (tx: Transaction) => Promise<any>,
+    Returns extends CB extends (tx: Transaction) => Promise<infer X> ? X : never,
+  >(query: CB): Promise<Returns> {
+    try {
+      if (isInZeroMutation()) {
+        const { tx } = mutatorContext()
+        return await query(tx)
+      }
+      // @ts-expect-error type
+      const output = await zeroDb.transaction(query, dummyTransactionInput)
+      return output
+    } catch (err) {
+      console.error(`Error running transaction(): ${err}`)
+      throw err
+    }
+  }
+
+  function query<R>(
+    cb: (q: Transaction['query']) => Query<any, Schema, R>
+  ): Promise<HumanReadable<R>> {
+    return transaction(async (tx) => {
+      return cb(tx.query)
+    }) as any
+  }
+
+  // register for global run() helper
+  if (queries) {
+    setCustomQueries(queries)
+  }
+
+  // server uses transaction-based execution
+  setRunner((queryObj) => {
+    return transaction(async (tx) => {
+      return tx.run(queryObj)
+    })
+  })
+
+  // This is needed temporarily and will be cleaned up in the future.
+  const dummyTransactionInput: TransactionProviderInput = {
+    clientGroupID: 'unused',
+    clientID: 'unused',
+    mutationID: 42,
+    upstreamSchema: 'unused',
+  }
+
+  return {
+    handleMutationRequest,
+    handleQueryRequest,
+    transaction,
+    mutate,
+    query,
+  }
+}

package/src/helpers/batchQuery.ts

@@ -0,0 +1,46 @@
+import { sleep } from '@take-out/helpers'
+
+import type { Query, Row } from '@rocicorp/zero'
+
+export async function batchQuery<Q extends Query<any, any, any>, Item extends Row<Q>>(
+  q: Q,
+  mapper: (items: Item[]) => Promise<void>,
+  {
+    chunk,
+    pause = 0,
+    stopAfter = 100_000,
+  }: {
+    chunk: number
+    pause?: number
+    stopAfter?: number
+  } = { chunk: 20 }
+) {
+  let hasMore = true
+  let last: Item | null = null
+  let iterations = 0
+
+  while (hasMore) {
+    let query = q.limit(chunk)
+
+    if (last) {
+      query = query.start(last)
+    }
+
+    const results = await query.run({ type: 'complete' })
+
+    await mapper(results as Item[])
+
+    if (results.length < chunk) {
+      hasMore = false
+    }
+
+    if (iterations > stopAfter) {
+      console.error(`[batchQuery] ‼️ stopping batch, ran ${stopAfter} chunks`)
+      break
+    }
+
+    if (pause) {
+      await sleep(pause)
+    }
+  }
+}

package/src/helpers/createMutators.ts

@@ -0,0 +1,138 @@
+import { mapObject, time } from '@take-out/helpers'
+
+import { runWithContext } from './mutatorContext'
+
+import { isBrowser, isServer } from '../constants'
+import type {
+  AuthData,
+  Can,
+  GenericModels,
+  GetZeroMutators,
+  MutatorContext,
+  Transaction,
+} from '../types'
+
+export function createMutators<Models extends GenericModels>({
+  environment,
+  authData,
+  createServerActions,
+  asyncTasks = [],
+  can,
+  models,
+}: {
+  environment: 'server' | 'client'
+  authData: AuthData | null
+  can: Can
+  models: Models
+  asyncTasks?: Array<() => Promise<void>>
+  createServerActions?: () => Record<string, any>
+}): GetZeroMutators<Models> {
+  const serverActions = createServerActions?.()
+
+  const modelMutators = mapObject(models, (val) => val.mutate) as {
+    [K in keyof typeof models]: (typeof models)[K]['mutate']
+  }
+
+  function withContext<Args extends any[]>(fn: (...args: Args) => Promise<void>) {
+    return async (tx: Transaction, ...args: Args): Promise<void> => {
+      const mutationContext: MutatorContext = {
+        tx,
+        authData,
+        environment,
+        can,
+        server:
+          environment === 'server'
+            ? ({
+                actions: serverActions || {},
+                asyncTasks: asyncTasks || {},
+              } as MutatorContext['server'])
+            : undefined,
+      }
+
+      return await runWithContext(mutationContext, () => {
+        // @ts-expect-error type shenanigan
+        // map to our mutations() helper
+        return fn(mutationContext, ...args)
+      })
+    }
+  }
+
+  function withDevelopmentLogging<Args extends any[]>(
+    name: string,
+    fn: (...args: Args) => Promise<void>
+  ) {
+    if (process.env.NODE_ENV !== 'development' && !process.env.IS_TESTING) {
+      return fn
+    }
+
+    return async (...args: Args): Promise<void> => {
+      const startTime = performance.now()
+
+      try {
+        if (isServer) {
+          console.info(`[mutator] ${name} start`)
+        }
+        const result = await fn(...args)
+        const duration = (performance.now() - startTime).toFixed(2)
+        if (isBrowser) {
+          console.groupCollapsed(`[mutator] ${name} completed in ${duration}ms`)
+          console.info('→', args[1])
+          console.info('←', result)
+          console.trace()
+          console.groupEnd()
+        } else {
+          // TODO in prod just track
+          console.info(`[mutator] ${name} completed in ${duration}ms`)
+        }
+        return result
+      } catch (error) {
+        const duration = (performance.now() - startTime).toFixed(2)
+        console.groupCollapsed(`[mutator] ${name} failed after ${duration}ms`)
+        console.error('error:', error)
+        console.info('arguments:', JSON.stringify(args[1], null, 2))
+        console.info('stack trace:', new Error().stack)
+        console.groupEnd()
+        throw error
+      }
+    }
+  }
+
+  function withTimeoutGuard<Args extends any[]>(
+    name: string,
+    fn: (...args: Args) => Promise<void>,
+    // don't want this too high - zero runs mutations in order and waits for the last to finish it seems
+    // so if one mutation gets stuck it will just sit there
+    timeoutMs: number = time.ms.minutes(1)
+  ) {
+    return async (...args: Args): Promise<void> => {
+      const timeoutPromise = new Promise<never>((_, reject) => {
+        setTimeout(() => {
+          reject(new Error(`[mutator] ${name} timeout after ${timeoutMs}ms`))
+        }, timeoutMs)
+      })
+
+      return Promise.race([fn(...args), timeoutPromise])
+    }
+  }
+
+  function decorateMutators<T extends Record<string, Record<string, any>>>(modules: T) {
+    const result: any = {}
+
+    for (const [moduleName, moduleExports] of Object.entries(modules)) {
+      result[moduleName] = {}
+      for (const [name, exportValue] of Object.entries(moduleExports)) {
+        if (typeof exportValue === 'function') {
+          const fullName = `${moduleName}.${name}`
+          result[moduleName][name] = withDevelopmentLogging(
+            fullName,
+            withTimeoutGuard(fullName, withContext(exportValue))
+          )
+        }
+      }
+    }
+
+    return result
+  }
+
+  return decorateMutators(modelMutators)
+}

package/src/helpers/didRunPermissionCheck.ts

@@ -0,0 +1,16 @@
+import { globalValue } from '@take-out/helpers'
+
+import type { MutatorContext } from '../types'
+
+const PermissionCheckRan = globalValue(
+  `on-zero:permissions-check`,
+  () => new WeakMap<MutatorContext, boolean>()
+)
+
+export const getDidRunPermissionCheck = (ctx: MutatorContext) => {
+  return PermissionCheckRan.get(ctx)
+}
+
+export const setDidRunPermissionCheck = (ctx: MutatorContext) => {
+  return PermissionCheckRan.set(ctx, true)
+}

package/src/helpers/ensureLoggedIn.ts

@@ -0,0 +1,11 @@
+import { ensure } from '@take-out/helpers'
+
+import { mutatorContext } from './mutatorContext'
+
+import type { AuthData } from '../types'
+
+export const ensureLoggedIn = (): AuthData => {
+  const { authData } = mutatorContext()
+  ensure(authData, 'logged in')
+  return authData
+}

package/src/helpers/getQueryOrMutatorAuthData.ts

@@ -0,0 +1,12 @@
+import { getAuthData } from '../state'
+import { isInZeroMutation, mutatorContext } from './mutatorContext'
+
+import type { AuthData } from '../types'
+
+export function getQueryOrMutatorAuthData(): AuthData | null {
+  if (isInZeroMutation()) {
+    return mutatorContext().authData as AuthData
+  }
+
+  return getAuthData()
+}

package/src/helpers/mutatorContext.ts

@@ -0,0 +1,25 @@
+import { createAsyncContext } from '@take-out/helpers'
+
+import type { MutatorContext } from '../types'
+
+const asyncContext = createAsyncContext<MutatorContext>()
+
+export function mutatorContext(): MutatorContext {
+  const currentContext = asyncContext.get()
+  if (!currentContext) {
+    throw new Error('mutatorContext must be called within a mutator')
+  }
+
+  return currentContext
+}
+
+export function isInZeroMutation() {
+  return !!asyncContext.get()
+}
+
+export function runWithContext<T>(
+  context: MutatorContext,
+  fn: () => T | Promise<T>
+): Promise<T> {
+  return asyncContext.run(context, fn)
+}

package/src/helpers/prettyFormatZeroQuery.ts

@@ -0,0 +1,167 @@
+import { ellipsis } from '@take-out/helpers'
+
+import type { Query } from '@rocicorp/zero'
+
+export const prettyFormatZeroQuery = (
+  query: Query<any, any, any>,
+  mode: 'full' | 'minimal' = 'full'
+): string => {
+  const astObject = query['_completeAst']?.()
+
+  if (!astObject) return ''
+
+  if (mode === 'minimal') {
+    return prettyFormatMinimal(astObject)
+  }
+  return prettyFormatFull(astObject)
+}
+
+const prettyFormatFull = (astObject: any, indent = 0): string => {
+  if (!astObject || !astObject.table) return ''
+
+  const spaces = '  '.repeat(indent)
+  let query = astObject.table
+  let hasChainedMethods = false
+
+  // Add where conditions
+  if (astObject.where) {
+    const whereClause = formatWhere(astObject.where)
+    if (hasChainedMethods) {
+      query += `\n${spaces}  ${whereClause}`
+    } else {
+      query += whereClause
+      hasChainedMethods = true
+    }
+  }
+
+  // Add limit
+  if (astObject.limit) {
+    const limitClause = `.limit(${astObject.limit})`
+    if (hasChainedMethods) {
+      query += `\n${spaces}  ${limitClause}`
+    } else {
+      query += limitClause
+      hasChainedMethods = true
+    }
+  }
+
+  // Add orderBy
+  if (astObject.orderBy && astObject.orderBy.length > 0) {
+    const orderClauses = astObject.orderBy
+      .map(([field, direction]: [string, string]) => `${field}, ${direction}`)
+      .join(', ')
+    const orderByClause = `.orderBy(${orderClauses})`
+    if (hasChainedMethods) {
+      query += `\n${spaces}  ${orderByClause}`
+    } else {
+      query += orderByClause
+      hasChainedMethods = true
+    }
+  }
+
+  // Add related queries
+  if (astObject.related && astObject.related.length > 0) {
+    astObject.related.forEach((rel: any) => {
+      if (rel.subquery) {
+        const alias = rel.subquery.alias || rel.subquery.table
+        const subQuery = prettyFormatFull(rel.subquery, indent + 1)
+        query += `\n${spaces}  .related(${alias}, q => q.${subQuery}`
+      }
+    })
+
+    // Add closing parentheses
+    const closingParens = ')'.repeat(astObject.related.length)
+    query += `\n${spaces}${closingParens}`
+  }
+
+  return query
+}
+
+const prettyFormatMinimal = (astObject: any): string => {
+  if (!astObject || !astObject.table) return ''
+
+  let query = astObject.table
+
+  // Add where conditions only
+  if (astObject.where) {
+    query += formatWhere(astObject.where).replace('.where(', '(')
+  }
+
+  // Add sub-queries info if present
+  if (astObject.related && astObject.related.length > 0) {
+    const subQueries = collectSubQueryTables(astObject.related)
+    const count = subQueries.length
+    const tableNames = subQueries.join(', ')
+    query += ` (+${count}: ${ellipsis(tableNames, 30)})`
+  }
+
+  return query
+}
+
+const collectSubQueryTables = (related: any[]): string[] => {
+  const tables: string[] = []
+
+  related.forEach((rel: any) => {
+    if (rel.subquery) {
+      const tableName = rel.subquery.alias || rel.subquery.table
+      tables.push(tableName)
+
+      // Recursively collect nested sub-queries
+      if (rel.subquery.related && rel.subquery.related.length > 0) {
+        tables.push(...collectSubQueryTables(rel.subquery.related))
+      }
+    }
+  })
+
+  return tables
+}
+
+const formatWhere = (where: any): string => {
+  if (!where) return ''
+
+  if (where.type === 'simple') {
+    const column = where.left?.name || where.left
+    const value = where.right?.value !== undefined ? where.right.value : where.right
+    const op = where.op || '='
+
+    // Special case: if column is "id" and op is "=" and value is a single item, show just the value
+    if (
+      column === 'id' &&
+      op === '=' &&
+      (typeof value === 'string' || typeof value === 'number')
+    ) {
+      return `(${value})`
+    }
+
+    if (op === '=') {
+      return `.where(${column}, ${value})`
+    }
+    return `.where(${column}, ${op}, ${value})`
+  }
+
+  if (where.type === 'and' && where.conditions) {
+    let result = ''
+    where.conditions.forEach((condition: any, index: number) => {
+      if (index === 0) {
+        result += formatWhere(condition)
+      } else {
+        result += `.and(${formatWhere(condition).slice(1)})` // Remove the leading dot
+      }
+    })
+    return result
+  }
+
+  if (where.type === 'or' && where.conditions) {
+    let result = ''
+    where.conditions.forEach((condition: any, index: number) => {
+      if (index === 0) {
+        result += formatWhere(condition)
+      } else {
+        result += `.or(${formatWhere(condition).slice(1)})` // Remove the leading dot
+      }
+    })
+    return result
+  }
+
+  return ''
+}