Using TanStack DB & ElectricSQL to build real-time apps with optimistic updates

// File path: apps/api/src/routes/main.ts
// A Hono-based backend route implementation

mainApp.post(
  "/shapes/auth-token",
  jwtAuth, // Middleware to validate the user's Auth Token
  zValidator("json", z.object({ table: z.string(), params: z.any() })),
  async (c) => {
    const { table, params } = c.req.valid("json");
    const user = c.get("user"); // User info parsed from the Auth Token

    let whereClause: string | undefined;

    // Dynamically generate the SQL WHERE clause based on the requested table and user identity
    if (table === "messages") {
      const conversationId = params?.conversationId;
      if (!conversationId) {
        // Disallow access to any messages if conversationId is not provided
        whereClause = `conversation_id = ''`;
      } else {
        // Business logic to verify user's ownership of the `conversationId` should be here.
        // For example, query the database to confirm the `conversationId` belongs to the current `user.userId`.
        const hasPermission = await checkConversationPermission(user.userId, conversationId);
        if (!hasPermission) {
          throw new ForbiddenError("No permission to access messages of this conversation");
        }

        // Authorize access to all messages in this conversation
        whereClause = `conversation_id = '${conversationId}'`;
      }
    } else {
      throw new ForbiddenError(`Access to table not allowed: ${table}`);
    }
    
    // Encode the authorization info (table name and WHERE clause) into a new Shape Token
    const shapeToken = await JWTService.generateShapeToken({ table, where: whereClause });

    return c.json({ token: shapeToken });
  }
);

// File path: apps/api/src/routes/main.ts
import { proxy } from 'hono/proxy'; // Ensure proxy is imported

mainApp.get(
  "/shape", 
  shapeAuth, // Key middleware: validates the Shape Token
  async (c) => {
    // 1. Get authorization parameters from the JWT payload.
    // The `shapeAuth` middleware has already validated the JWT and placed its payload
    // in the request context. This is the source of truth, not the client's URL query params.
    const { table, where } = c.get("shapePayload");
    
    const electricServiceUrl = new URL(`${process.env.ELECTRIC_API_URL}/v1/shape`);

    // 2. Forward only non-security-related client query parameters (e.g., `offset`).
    const clientRequestUrl = new URL(c.req.url);
    clientRequestUrl.searchParams.forEach((value, key) => {
      if (key !== "table" && key !== "where") {
        electricServiceUrl.searchParams.set(key, value);
      }
    });

    // 3. Construct the final request using the validated parameters from the JWT.
    // This prevents the client from tampering with the request to access unauthorized tables or rows.
    electricServiceUrl.searchParams.set("table", table);
    if (where) {
      electricServiceUrl.searchParams.set("where", where);
    }

    // 4. Proxy the well-formed, secure request to the ElectricSQL service.
    // The proxy function is from the 'hono/proxy' library.
    return proxy(electricServiceUrl.toString(), {
      ...c.req,
      headers: {
        ...c.req.headers,
      },
    });
});

// File path: apps/web/src/hooks/data/collections.ts
import { createCollection } from '@tanstack/react-db';
import { electricCollectionOptions } from '@tanstack/electric-db-collection';
import { selectMessageSchema, type Message } from '@/db/schema'; // Zod Schema
import { createShapeConfig } from './shape-config'; // Import the Shape config factory
import { getApiBaseUrl } from './api'; // Import the API base URL getter

// Internal function to create the actual collection
function _createMessagesCollection(conversationId: string) {
  // Optional: specify columns to sync for performance optimization
  const columnsToSync = [
    'id',
    'conversation_id',
    'content',
    'role',
    'status',
    'created_at',
    'updated_at',
  ].join(',');

  // `createShapeConfig` is a helper function encapsulating auth and error handling logic
  const shapeConfig = createShapeConfig('messages', { conversationId }, undefined, columnsToSync);

  return createCollection(
    electricCollectionOptions({
      id: `messages-${conversationId}`,
      schema: selectMessageSchema,
      shapeOptions: {
        url: shapeConfig.url || `${getApiBaseUrl()}/shape`,
        params: shapeConfig.params,
        headers: shapeConfig.headers,
        onError: shapeConfig.onError,
        parser: shapeConfig.parser, // Optional custom parser for data transformation
      },
      getKey: (item: Message) => item.id,
    }),
  );
}

// Type alias for better type safety
type MessagesCollectionType = ReturnType<typeof _createMessagesCollection>;

// Global cache to avoid creating duplicate Collection instances
const messagesCollectionsCache = new Map<string, MessagesCollectionType>();

/**
 * Creates a messages collection with caching
 *
 * @example
 * // ✅ Correct usage
 * const messagesCollection = createMessagesCollection(conversationId);
 *
 * // ❌ Wrong usage - DO NOT wrap in useMemo
 * const messagesCollection = useMemo(() => createMessagesCollection(conversationId), [conversationId]);
 *
 * @param conversationId - The conversation ID to filter messages
 * @returns Cached collection instance
 */
export function createMessagesCollection(conversationId: string) {
  const cacheKey = `messages-${conversationId}`;

  if (messagesCollectionsCache.has(cacheKey)) {
    return messagesCollectionsCache.get(cacheKey)!;
  }

  const collection = _createMessagesCollection(conversationId);
  messagesCollectionsCache.set(cacheKey, collection);
  return collection;
}

// File path: apps/web/src/lib/api.ts

// In a real application, this URL might be dynamic (e.g., from .env or a Zustand store)
export function getApiBaseUrl(): string {
  return 'https://your-api.example.com'; 
}

// In a real application, API call functions would use fetch or axios
export async function sendQuestionApi(params: any): Promise<{ txid: string }> {
  const response = await fetch(`${getApiBaseUrl()}/questions/send`, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${getAuthToken()}`, // Function to get the user's Auth Token
    },
    body: JSON.stringify(params),
  });
  if (!response.ok) throw new Error('API request failed');
  return response.json();
}

// File path: apps/web/src/stores/shape-token-store.ts
// A Shape Token cache implemented with Zustand

import { create } from 'zustand';

interface ShapeTokenState {
  tokenCache: Map<string, string>; // Key: `${table}|${JSON.stringify(params)}`
  getToken: (key: string) => string | undefined;
  setToken: (key: string, token: string) => void;
  removeToken: (key: string) => void;
}

export const useShapeTokenStore = create<ShapeTokenState>((set, get) => ({
  tokenCache: new Map(),
  getToken: key => get().tokenCache.get(key),
  setToken: (key, token) => set(state => ({ tokenCache: new Map(state.tokenCache).set(key, token) })),
  removeToken: (key) => {
    set((state) => {
      const newCache = new Map(state.tokenCache);
      newCache.delete(key);
      return { tokenCache: newCache };
    });
  },
}));

// File path: apps/web/src/hooks/data/shape-config.ts
import { type ShapeOptions } from '@electric-sql/react';
import { FetchError } from '@electric-sql/client';
import { useShapeTokenStore } from '@/stores/shape-token-store';
import { getApiBaseUrl } from '@/lib/api';

// A simplified fetchShapeToken function
async function fetchShapeToken(table: string, params: Record<string, any>): Promise<string> {
   const response = await fetch(`${getApiBaseUrl()}/shapes/auth-token`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${getAuthToken()}`, // Function to get the user's Auth Token
      },
      body: JSON.stringify({ table, params }),
    });
    if (!response.ok) throw new Error('Failed to fetch shape token');
    const data = await response.json();
    return data.token;
}


export function createShapeConfig<T extends Row>(
  table: string,
  params: Record<string, any>,
): ShapeOptions<T> {
  const cacheKey = `${table}|${JSON.stringify(params)}`;

  const getAuthHeader = async () => {
    // Check cache first
    const cachedToken = useShapeTokenStore.getState().getToken(cacheKey);
    if (cachedToken) return `Bearer ${cachedToken}`;

    // On cache miss, fetch a new token
    const newShapeToken = await fetchShapeToken(table, params);
    useShapeTokenStore.getState().setToken(cacheKey, newShapeToken);
    return `Bearer ${newShapeToken}`;
  };

  const onErrorFunc = async (error: Error) => {
    // Handle only 401 errors, which usually mean an expired token
    if (error instanceof FetchError && error.status === 401) {
      console.warn(`[Shape Auth] 401 Unauthorized. Invalidating token and retrying...`);
      // Key step: remove the expired token from the cache
      useShapeTokenStore.getState().removeToken(cacheKey);
      // Return an empty object to signal to ElectricSQL that this is a recoverable error, and it should retry automatically
      return {}; 
    }
    throw error;
  };

  return {
    params: { table, ...params },
    headers: { Authorization: getAuthHeader },
    onError: onErrorFunc,
  };
}

// File path: apps/web/src/hooks/data/actions.ts
import { createTransaction } from '@tanstack/react-db';
import { createMessagesCollection } from './collections';
import { sendQuestionApi } from '@/lib/api'; // API call function
import { createQuestionPlaceholder, createAnswerPlaceholder } from './placeholders'; // Placeholder factories

export async function sendQuestionAction(params: {
  conversationId: string;
  questionContent: string;
  questionId: string;  // Pre-generated on the client
  answerId: string;    // Pre-generated on the client
}) {
  const { conversationId, questionContent, questionId, answerId } = params;
  const messagesCollection = createMessagesCollection(conversationId);

  const transaction = createTransaction({
    autoCommit: false,
    mutationFn: async () => {
      // 1. Call the backend API
      const response = await sendQuestionApi({
        conversationId,
        questionId,
        answerId,
        content: questionContent,
      });
      if (!response.txid) throw new Error('API did not return a transaction ID');

      // 2. Wait for ElectricSQL to sync this transaction back to the client.
      // This is the key to the seamless "optimistic -> real data" transition.
      // IMPORTANT: awaitTxId must be awaited, otherwise optimistic updates will be ignored.
      await messagesCollection.utils.awaitTxId(response.txid);
    },
  });

  // Optimistic Update
  // This part executes synchronously and immediately, making the UI feel instant.
  transaction.mutate(() => {
    const questionMessage = createQuestionPlaceholder({
      id: questionId,
      conversation_id: conversationId,
      content: questionContent,
    });
    const answerMessage = createAnswerPlaceholder({
      id: answerId,
      conversation_id: conversationId,
    });
    messagesCollection.insert([questionMessage, answerMessage]);
  });

  // Commit the transaction, which triggers the execution of `mutationFn`
  await transaction.commit();
}

// File path: apps/web/src/hooks/data/placeholders.ts
import type { Message } from '@repo/shared/types';

export function createQuestionPlaceholder(data: Partial<Message>): Message {
  const now = new Date();
  return {
    id: data.id!,
    conversation_id: data.conversation_id!,
    content: data.content || '',
    role: 'user', // or 'question'
    status: 'completed',
    created_at: now,
    updated_at: now,
    // ... other schema-required fields should have default values
  };
}

export function createAnswerPlaceholder(data: Partial<Message>): Message {
  const now = new Date();
  return {
    id: data.id!,
    conversation_id: data.conversation_id!,
    content: 'Generating answer...',
    role: 'assistant', // or 'answer'
    status: 'pending',
    created_at: new Date(now.getTime() + 1), // Ensure sequence
    updated_at: new Date(now.getTime() + 1),
    // ... other schema-required fields should have default values
  };
}

// File path: apps/api/src/routes/main.ts
mainApp.post("/questions/send", jwtAuth, async (c) => {
  const { conversationId, questionId, answerId, content } = c.req.valid("json");
  const userId = c.get("user").userId;

  // 1. Create the question and answer placeholder in a single transaction
  const { txid } = await messageService.createQuestionAndAnswerPlaceholder({
    conversationId,
    userId,
    questionId,
    questionContent: content,
    answerId,
  });

  // 2. Asynchronously trigger answer generation
  setImmediate(() => {
    generateAnswer(conversationId, answerId, content)
      .catch(err => console.error("Answer generation failed:", err));
  });
  
  // 3. Immediately return the transaction ID
  return c.json({ success: true, txid });
});

// File path: apps/api/src/services/message.service.ts

export class MessageService {
  // ... other methods

  /**
   * A transaction wrapper that ensures all operations run within a transaction and returns the transaction ID.
   */
  private async withTransaction<T>(
    callback: (tx: DrizzleTransaction, txid: number) => Promise<T>,
  ): Promise<T> {
    // Use Drizzle's transaction method to start a database transaction
    return await db.transaction(async (tx) => {
      // Inside the transaction, first get the current transaction ID
      const txid = await this.getTxId(tx);
      // Then, execute the callback containing all database operations
      return await callback(tx, txid);
    });
  }

  /**
   * Gets the current transaction's ID.
   * Internal function used by various transaction wrappers
   *
   * Uses pg_current_xact_id()::xid::text to get the 32-bit internal xid,
   * which matches the ID sent in PostgreSQL logical replication streams,
   * and is the ID we use for matching in Electric
   */
  private async getTxId(tx: DrizzleTransaction): Promise<number> {
    try {
      // Use pg_current_xact_id()::xid::text to get the 32-bit internal xid
      // ::xid conversion removes the epoch, giving the raw 32-bit value
      // This matches the value PostgreSQL sends in logical replication streams
      const result = await tx.execute(sql`SELECT pg_current_xact_id()::xid::text as txid`);
      const txid = (result.rows[0] as { txid: string })?.txid;

      if (txid === undefined) {
        throw new Error("Failed to get transaction ID from pg_current_xact_id");
      }

      return Number.parseInt(txid, 10);
    }
    catch (error) {
      // Fallback to deprecated txid_current()
      console.warn("[Transaction] pg_current_xact_id() failed, falling back to txid_current():", error);
      try {
        const result = await tx.execute(sql`SELECT txid_current() as txid`);
        const txid = (result.rows[0] as { txid: bigint })?.txid;

        if (txid === undefined) {
          throw new Error("Failed to get transaction ID from txid_current");
        }

        return Number(txid);
      }
      catch (fallbackError) {
        console.error("[Transaction] Both txid methods failed:", fallbackError);
        throw new Error("Failed to get transaction ID from all available methods");
      }
    }
  }

  /**
   * Creates the question and answer placeholders in a single transaction.
   */
  async createQuestionAndAnswerPlaceholder(params: ...): Promise<{ txid: number }> {
    // Execute operations via withTransaction to ensure atomicity and get the txid
    return this.withTransaction(async (tx, txid) => {
      // 1. Insert question message...
      await tx.insert(messages).values({ ... });

      // 2. Insert answer placeholder...
      await tx.insert(messages).values({ ... });
      
      // 3. Return the txid obtained from the wrapper
      return { txid };
    });
  }
}

// File path: apps/web/src/components/question-answer-list.tsx
import { useLiveQuery } from '@tanstack/react-db';
import { createMessagesCollection } from '@/hooks/data/collections';

function QuestionAnswerList({ conversationId }: { conversationId: string }) {
  // Subscribe to the message collection for a specific conversation
  const messagesCollection = useMemo(() => createMessagesCollection(conversationId), [conversationId]);

  // The new useLiveQuery syntax with destructured data and loading state
  const {
    data: messages,
    isLoading,
  } = useLiveQuery(q =>
    q.from({ messagesCollection })
      .orderBy(({ messagesCollection }) => messagesCollection.created_at, 'desc')
);

  // Handle loading state - when there's no data, it stays in loading state
  if (isLoading) {
    return <div className="loading">Loading messages...</div>;
  }

  return (
    <div className="message-container">
      {messages?.map(msg => (
        <div key={msg.id} className={`message role-${msg.role} status-${msg.status}`}>
          <p>{msg.content}</p>
        </div>
      ))}
    </div>
  );
}

// Database schema definition with Drizzle ORM
const messages = pgTable("messages", {
  id: text("id").primaryKey(),
  content: text("content").notNull(),
  metadata: jsonb("metadata").$type<{
    author: { name: string; id: string };
    tags: string[];
    priority: number;
  }>(),
});

// Type inference from Drizzle
type Message = typeof messages.$inferSelect;
// Message.metadata type is: { author: ..., tags: ..., priority: ... } | null

// What ElectricSQL returns after synchronization
const selectMessageSchema = createSelectSchema(messages);
type ElectricMessage = z.infer<typeof selectMessageSchema>;
// ElectricMessage.metadata type is: Json | null (type information lost!)

// Using Drizzle-generated types and Zod schemas
const messages = pgTable("messages", {
  id: text("id").primaryKey(),
  content: text("content").notNull(),
  metadata: jsonb("metadata").$type<MessageMetadata>(),
  log_data: jsonb("log_data").$type<LogType1 | LogType2 | LogType3>(),
});

type Message = typeof messages.$inferSelect;
const selectMessageSchema = createSelectSchema(messages);

// Transform function to restore proper types after ElectricSQL sync
function transformElectricMessagesToTyped(
  electricMessages: Array<z.infer<typeof selectMessageSchema>>
): Message[] {
  return electricMessages.map(msg => ({
    ...msg,
    // Type assertions required for JSONB fields
    metadata: msg.metadata as Message['metadata'],
    log_data: msg.log_data as Message['log_data'],
  }));
}

// Usage in collection with automatic transformation
export function createMessagesCollection(conversationId: string) {
  const collection = createCollection(
    electricCollectionOptions({
      id: `messages-${conversationId}`,
      schema: selectMessageSchema,
      shapeOptions: { /* ... */ },
      getKey: item => item.id,
    }),
  );

  // Wrap collection methods to include transformation
  return {
    ...collection,
    findMany: async (filter?: any) => {
      const results = await collection.findMany(filter);
      return transformElectricMessagesToTyped(results);
    },
    // Add other wrapped methods as needed
  };
}

// Real-world example: Collection factory with proper caching
import { pgTable, text, jsonb, timestamp } from 'drizzle-orm/pg-core';
import { createSelectSchema } from 'drizzle-zod';
import type { ShapeOptions } from '@electric-sql/react';

// Define your table with typed JSONB
const documents = pgTable("documents", {
  id: text("id").primaryKey(),
  status: text("status").notNull(),
  metadata: jsonb("metadata").$type<{
    processor: string;
    version: string;
    tags: string[];
  }>(),
  created_at: timestamp("created_at").defaultNow(),
  updated_at: timestamp("updated_at").defaultNow(),
});

// Get proper TypeScript types
type Document = typeof documents.$inferSelect;
const selectDocumentSchema = createSelectSchema(documents);

// Internal collection factory (not exported)
function _createDocumentsCollection() {
  // Optimize by selecting specific columns
  const columnsToSync = [
    'id',
    'status',
    'metadata',
    'created_at',
    'updated_at',
  ].join(',');

  // Custom parser for date handling and type restoration
  const parser = {
    timestamp: (date: string) => new Date(date),
  } as unknown as ShapeOptions['parser'];

  const shapeConfig = createShapeConfig('documents', {}, parser, columnsToSync);

  return createCollection(
    electricCollectionOptions({
      id: 'documents',
      schema: selectDocumentSchema,
      shapeOptions: {
        url: shapeConfig.url,
        params: shapeConfig.params,
        headers: shapeConfig.headers,
        onError: shapeConfig.onError,
        parser: shapeConfig.parser,
      },
      getKey: (item: Document) => item.id,
    }),
  );
}

// Type alias for the collection
type DocumentsCollectionType = ReturnType<typeof _createDocumentsCollection>;

// Cache instance
const documentsCollectionsCache = new Map<string, DocumentsCollectionType>();

/**
 * Creates a documents collection with caching
 *
 * @example
 * // ✅ Correct usage - direct call
 * const docsCollection = createDocumentsCollection();
 *
 * // ❌ Wrong usage - DO NOT wrap in useMemo
 * const docsCollection = useMemo(() => createDocumentsCollection(), []);
 *
 * @returns Cached collection instance
 */
export function createDocumentsCollection() {
  const cacheKey = 'documents';

  if (documentsCollectionsCache.has(cacheKey)) {
    return documentsCollectionsCache.get(cacheKey)!;
  }

  const collection = _createDocumentsCollection();
  documentsCollectionsCache.set(cacheKey, collection);
  return collection;
}

// Multiple database queries to check permissions
const hasPermission = await checkConversationPermission(user.userId, conversationId);
if (!hasPermission) {
  throw new ForbiddenError("No permission");
}
whereClause = `conversation_id = '${conversationId}'`;

// Single WHERE clause with subquery handles everything
whereClause = `conversation_id IN (
  SELECT c.id FROM conversations c
  JOIN conversation_members cm ON c.id = cm.conversation_id
  WHERE cm.user_id = '${user.userId}'
    AND cm.role IN ('owner', 'member')
)`;