drizzleConvex.ts

import { drizzle } from 'drizzle-orm/node-postgres'
import { desc, asc, eq } from 'drizzle-orm'
import { numbersTable } from '~/db/schema'

// Type definitions to match Convex patterns
type TableName = 'numbers'
type OrderDirection = 'asc' | 'desc'

// Table mapping - extend this as you add more tables
const tableMap = {
  numbers: numbersTable,
} as const

// Type mapping for table results
type TableResult<T extends TableName> = T extends 'numbers'
  ? {
      _id: number
      createdAt: Date | null
      value: string[]
    }
  : never

// Type mapping for table IDs
type TableId<T extends TableName> = T extends 'numbers' ? number : never

// Type mapping for insert data
type TableInsertData<T extends TableName> = T extends 'numbers'
  ? { value: number }
  : never

// Query builder class that mimics Convex's chainable API
class QueryBuilder<T> {
  private db: ReturnType<typeof drizzle>
  private table: any
  private orderDirection: OrderDirection = 'asc'
  private limitCount?: number
  private whereConditions: any[] = []

  constructor(db: ReturnType<typeof drizzle>, table: any) {
    this.db = db
    this.table = table
  }

  // Mimic Convex's order() method
  order(direction: OrderDirection): this {
    this.orderDirection = direction
    return this
  }

  // Mimic Convex's take() method
  take(count: number): this {
    this.limitCount = count
    return this
  }

  // Mimic Convex's filter() method (simplified version)
  filter(condition: any): this {
    this.whereConditions.push(condition)
    return this
  }

  // Execute the query and return results
  async execute(): Promise<T[]> {
    // Build the query step by step to avoid TypeScript issues
    let queryBuilder = this.db.select().from(this.table)

    // Apply where conditions
    if (this.whereConditions.length > 0) {
      queryBuilder = queryBuilder.where(this.whereConditions[0]) as any
    }

    // Apply ordering
    if (this.orderDirection === 'desc') {
      queryBuilder = queryBuilder.orderBy(desc(this.table.createdAt)) as any
    } else {
      queryBuilder = queryBuilder.orderBy(asc(this.table.createdAt)) as any
    }

    // Apply limit
    if (this.limitCount) {
      queryBuilder = queryBuilder.limit(this.limitCount) as any
    }

    return await queryBuilder
  }
}

// Database context class that mimics Convex's ctx.db
class DrizzleConvexDB {
  private db: ReturnType<typeof drizzle>

  constructor(databaseUrl: string) {
    this.db = drizzle(databaseUrl)
  }

  // Mimic Convex's ctx.db.query() method
  query<T extends TableName>(
    tableName: T,
  ): QueryBuilder<TableResult<T>> & Promise<TableResult<T>[]> {
    const table = tableMap[tableName]
    const builder = new QueryBuilder<TableResult<T>>(this.db, table)

    // Make the builder thenable so it can be awaited directly
    const thenableBuilder = builder as QueryBuilder<TableResult<T>> &
      Promise<TableResult<T>[]>
    thenableBuilder.then = (onfulfilled?: any, onrejected?: any) => {
      return builder.execute().then(onfulfilled, onrejected)
    }
    thenableBuilder.catch = (onrejected?: any) => {
      return builder.execute().catch(onrejected)
    }
    thenableBuilder.finally = (onfinally?: any) => {
      return builder.execute().finally(onfinally)
    }

    return thenableBuilder
  }

  // Mimic Convex's ctx.db.insert() method
  async insert<T extends TableName>(
    tableName: T,
    data: TableInsertData<T>,
  ): Promise<TableId<T>> {
    const table = tableMap[tableName]

    if (tableName === 'numbers') {
      // Handle the specific case for numbers table
      const transformedData = {
        value: Array.isArray((data as any).value)
          ? (data as any).value.map((v: number) => v.toString())
          : [(data as any).value.toString()],
      }
      const result = await this.db
        .insert(table as any)
        .values(transformedData as any)
        .returning({ id: table._id })
      return result[0]?.id as unknown as TableId<T>
    }

    // For other tables, insert as-is
    const result = await this.db
      .insert(table as any)
      .values(data as any)
      .returning()
    return (result as any)[0] as unknown as TableId<T>
  }

  // Mimic Convex's ctx.db.get() method
  async get(id: number): Promise<any | null> {
    // This is a simplified version - in practice you'd need to determine the table from the ID
    const result = await this.db
      .select()
      .from(numbersTable)
      .where(eq(numbersTable._id, id))
      .limit(1)
    return result[0] || null
  }

  // Mimic Convex's ctx.db.patch() method
  async patch(id: number, updates: any): Promise<void> {
    // Simplified - you'd need table detection logic
    await this.db
      .update(numbersTable)
      .set(updates)
      .where(eq(numbersTable._id, id))
  }

  // Mimic Convex's ctx.db.delete() method
  async delete(id: number): Promise<void> {
    // Simplified - you'd need table detection logic
    await this.db.delete(numbersTable).where(eq(numbersTable._id, id))
  }

  // Add transaction support like Drizzle
  async transaction<T>(
    callback: (tx: DrizzleConvexDB) => Promise<T>,
  ): Promise<T> {
    return await this.db.transaction(async (tx) => {
      // Create a new DrizzleConvexDB instance with the transaction
      const transactionalDB = new DrizzleConvexDB('')
      transactionalDB.db = tx as any
      return await callback(transactionalDB)
    })
  }
}

// Context creator function
export function createDrizzleConvexContext(databaseUrl: string) {
  return {
    db: new DrizzleConvexDB(databaseUrl),
  }
}

// Export the context type for TypeScript
export type DrizzleConvexContext = ReturnType<typeof createDrizzleConvexContext>

tanstackStartConvex.ts

import { createServerFn } from '@tanstack/react-start'
import { z } from 'zod'
import {
  createDrizzleConvexContext,
  type DrizzleConvexContext,
} from './drizzleConvex'

// Convex-like validator system using Zod
export const v = {
  number: () => z.number(),
  string: () => z.string(),
  boolean: () => z.boolean(),
  array: <T extends z.ZodTypeAny>(schema: T) => z.array(schema),
  object: <T extends z.ZodRawShape>(shape: T) => z.object(shape),
  optional: <T extends z.ZodTypeAny>(schema: T) => schema.optional(),
  null: () => z.null(),
  union: <T extends readonly [z.ZodTypeAny, ...z.ZodTypeAny[]]>(
    ...schemas: T
  ) => z.union(schemas),
  literal: <T extends string | number | boolean>(value: T) => z.literal(value),
  id: (tableName: string) => z.number(), // Simplified - in real Convex this would be a branded type
}

// Type for our Convex-like function configs
type ConvexFunctionConfig<TReturn = any> = {
  validator: (input: any) => any
  handler: ({ data }: { data: any }) => Promise<TReturn>
}

// Enhanced context type with runQuery and runMutation like Convex
type ConvexLikeContext = DrizzleConvexContext & {
  runQuery: <TReturn>(
    config: ConvexFunctionConfig<TReturn>,
    args: any,
  ) => Promise<TReturn>
  runMutation: <TReturn>(
    config: ConvexFunctionConfig<TReturn>,
    args: any,
  ) => Promise<TReturn>
}

// Action context type without direct database access (like Convex)
type ConvexActionContext = {
  runQuery: <TReturn>(
    config: ConvexFunctionConfig<TReturn>,
    args: any,
  ) => Promise<TReturn>
  runMutation: <TReturn>(
    config: ConvexFunctionConfig<TReturn>,
    args: any,
  ) => Promise<TReturn>
}

// Create database context singleton
let dbContext: DrizzleConvexContext | null = null

function getDbContext(): ConvexLikeContext {
  if (!dbContext) {
    dbContext = createDrizzleConvexContext(process.env.DATABASE_URL!)
  }

  // Add runQuery and runMutation to the context like Convex
  return {
    ...dbContext,
    runQuery: async <TReturn>(
      config: ConvexFunctionConfig<TReturn>,
      args: any,
    ): Promise<TReturn> => {
      // Validate the arguments
      const validatedArgs = config.validator(args)
      // Call the handler directly with the validated args
      return await config.handler({ data: validatedArgs })
    },
    runMutation: async <TReturn>(
      config: ConvexFunctionConfig<TReturn>,
      args: any,
    ): Promise<TReturn> => {
      // Validate the arguments
      const validatedArgs = config.validator(args)
      // Call the handler directly with the validated args
      return await config.handler({ data: validatedArgs })
    },
  }
}

// Helper function to create a Convex-like query
export function query<
  TArgs extends Record<string, z.ZodTypeAny>,
  TReturn,
>(config: {
  args: TArgs
  returns?: z.ZodTypeAny
  handler: (
    ctx: ConvexLikeContext,
    args: { [K in keyof TArgs]: z.infer<TArgs[K]> },
  ) => Promise<TReturn>
}) {
  // Return the createServerFn configuration that you can use directly
  return {
    validator: (input: any) => {
      const argsSchema = z.object(config.args)
      return argsSchema.parse(input)
    },
    handler: async ({ data }: { data: any }): Promise<TReturn> => {
      const baseCtx = getDbContext()
      // Run the query in a transaction for consistency
      return await baseCtx.db.transaction(async (tx) => {
        // Create a new context with the transaction
        const transactionalCtx = {
          ...baseCtx,
          db: tx,
        }
        return await config.handler(transactionalCtx, data)
      })
    },
  }
}

// Helper function to create a Convex-like mutation
export function mutation<
  TArgs extends Record<string, z.ZodTypeAny>,
  TReturn,
>(config: {
  args: TArgs
  returns?: z.ZodTypeAny
  handler: (
    ctx: ConvexLikeContext,
    args: { [K in keyof TArgs]: z.infer<TArgs[K]> },
  ) => Promise<TReturn>
}) {
  // Return the createServerFn configuration that you can use directly
  return {
    validator: (input: any) => {
      const argsSchema = z.object(config.args)
      return argsSchema.parse(input)
    },
    handler: async ({ data }: { data: any }): Promise<TReturn> => {
      const baseCtx = getDbContext()
      // Run the mutation in a transaction for ACID properties
      return await baseCtx.db.transaction(async (tx) => {
        // Create a new context with the transaction
        const transactionalCtx = {
          ...baseCtx,
          db: tx,
        }
        return await config.handler(transactionalCtx, data)
      })
    },
  }
}

// Helper function to create a Convex-like action
export function action<
  TArgs extends Record<string, z.ZodTypeAny>,
  TReturn,
>(config: {
  args: TArgs
  returns?: z.ZodTypeAny
  handler: (
    ctx: ConvexActionContext,
    args: { [K in keyof TArgs]: z.infer<TArgs[K]> },
  ) => Promise<TReturn>
}) {
  // Return the createServerFn configuration that you can use directly
  return {
    validator: (input: any) => {
      const argsSchema = z.object(config.args)
      return argsSchema.parse(input)
    },
    handler: async ({ data }: { data: any }): Promise<TReturn> => {
      // Actions don't have direct database access, only runQuery and runMutation
      const actionCtx: ConvexActionContext = {
        runQuery: async <TReturn>(
          config: ConvexFunctionConfig<TReturn>,
          args: any,
        ): Promise<TReturn> => {
          // Validate the arguments
          const validatedArgs = config.validator(args)
          // Call the handler directly with the validated args
          return await config.handler({ data: validatedArgs })
        },
        runMutation: async <TReturn>(
          config: ConvexFunctionConfig<TReturn>,
          args: any,
        ): Promise<TReturn> => {
          // Validate the arguments
          const validatedArgs = config.validator(args)
          // Call the handler directly with the validated args
          return await config.handler({ data: validatedArgs })
        },
      }
      return await config.handler(actionCtx, data)
    },
  }
}

// Helper for calling other functions (mimics ctx.runQuery, ctx.runMutation)
export async function runQuery<T>(fn: any, args: any): Promise<T> {
  // This is a simplified version - in practice you'd need more sophisticated function calling
  return await fn({ data: args })
}

export async function runMutation<T>(fn: any, args: any): Promise<T> {
  // This is a simplified version - in practice you'd need more sophisticated function calling
  return await fn({ data: args })
}

// Keep the old runFunction for backward compatibility
export async function runFunction<T>(fn: any, args: any): Promise<T> {
  return await fn({ data: args })
}