@dataql/react-native/src/wal/SQLiteWAL.ts

DataQL React Native SDK with offline-first capabilities and clean API

import {
  WALEntry,
  WALTransaction,
  WALRecoveryPoint,
  WALConfig,
} from "../../../core/src/lib/wal/types";
import { BaseWAL } from "../../../core/src/lib/wal/BaseWAL";
import { DrizzleD1Database } from "drizzle-orm/d1";
import { sqliteTable, text, integer } from "drizzle-orm/sqlite-core";
import { eq, and, gte, lt, sql } from "drizzle-orm";

// WAL schema for SQLite
export const walEntries = sqliteTable("wal_entries", {
  id: text("id").primaryKey(),
  walId: text("wal_id").notNull().unique(),
  sessionId: text("session_id").notNull(),
  collection: text("collection").notNull(),
  operation: text("operation").notNull(),
  data: text("data").notNull(), // JSON string
  previousData: text("previous_data"), // JSON string
  transactionId: text("transaction_id"),
  status: text("status").notNull(),
  checksum: text("checksum"),
  metadata: text("metadata"), // JSON string
  timestamp: integer("timestamp").notNull(),
  createdAt: integer("created_at").notNull(),
  updatedAt: integer("updated_at"),
});

export const walTransactions = sqliteTable("wal_transactions", {
  id: text("id").primaryKey(),
  walId: text("wal_id").notNull().unique(),
  sessionId: text("session_id").notNull(),
  status: text("status").notNull(),
  operationCount: integer("operation_count").notNull(),
  timestamp: integer("timestamp").notNull(),
  createdAt: integer("created_at").notNull(),
  updatedAt: integer("updated_at"),
});

export const walRecoveryPoints = sqliteTable("wal_recovery_points", {
  id: text("id").primaryKey(),
  walId: text("wal_id").notNull().unique(),
  sessionId: text("session_id").notNull(),
  walEntryId: text("wal_entry_id").notNull(),
  timestamp: integer("timestamp").notNull(),
  createdAt: integer("created_at").notNull(),
});

export class SQLiteWAL extends BaseWAL {
  private db: DrizzleD1Database<any>;

  constructor(db: DrizzleD1Database<any>, config: Partial<WALConfig> = {}) {
    super(config);
    this.db = db;
    this.initializeSchema();
  }

  private async initializeSchema(): Promise<void> {
    try {
      // Create indexes for better performance
      await this.db.run(`
        CREATE INDEX IF NOT EXISTS idx_wal_entries_session_timestamp 
        ON wal_entries(session_id, timestamp)
      `);

      await this.db.run(`
        CREATE INDEX IF NOT EXISTS idx_wal_entries_status_timestamp 
        ON wal_entries(status, timestamp)
      `);

      await this.db.run(`
        CREATE INDEX IF NOT EXISTS idx_wal_entries_transaction_id 
        ON wal_entries(transaction_id) WHERE transaction_id IS NOT NULL
      `);

      await this.db.run(`
        CREATE INDEX IF NOT EXISTS idx_wal_transactions_session_timestamp 
        ON wal_transactions(session_id, timestamp)
      `);

      await this.db.run(`
        CREATE INDEX IF NOT EXISTS idx_wal_recovery_session_timestamp 
        ON wal_recovery_points(session_id, timestamp)
      `);
    } catch (error) {
      console.warn("Failed to create WAL indexes:", error);
    }
  }

  protected async persistEntry(entry: WALEntry): Promise<void> {
    try {
      const now = Date.now();

      await this.db.insert(walEntries).values({
        id: `sqlite_${Date.now()}_${Math.random().toString(36).substring(2)}`,
        walId: entry.id,
        sessionId: entry.sessionId,
        collection: entry.collection,
        operation: entry.operation,
        data: JSON.stringify(entry.data),
        previousData: entry.previousData
          ? JSON.stringify(entry.previousData)
          : null,
        transactionId: entry.transactionId || null,
        status: entry.status,
        checksum: entry.checksum || null,
        metadata: entry.metadata ? JSON.stringify(entry.metadata) : null,
        timestamp: entry.timestamp.getTime(),
        createdAt: now,
        updatedAt: null,
      });
    } catch (error) {
      throw new Error(
        `Failed to persist WAL entry: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async persistTransaction(
    transaction: WALTransaction
  ): Promise<void> {
    try {
      const now = Date.now();

      // Begin transaction in SQLite
      await this.db.run("BEGIN TRANSACTION");

      try {
        // Insert transaction record
        await this.db.insert(walTransactions).values({
          id: `sqlite_txn_${Date.now()}_${Math.random().toString(36).substring(2)}`,
          walId: transaction.id,
          sessionId: transaction.sessionId,
          status: transaction.status,
          operationCount: transaction.operations.length,
          timestamp: transaction.timestamp.getTime(),
          createdAt: now,
          updatedAt: null,
        });

        // Insert all operations
        for (const operation of transaction.operations) {
          await this.persistEntry(operation);
        }

        await this.db.run("COMMIT");
      } catch (error) {
        await this.db.run("ROLLBACK");
        throw error;
      }
    } catch (error) {
      throw new Error(
        `Failed to persist WAL transaction: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async updateEntryStatus(
    entryId: string,
    status: WALEntry["status"]
  ): Promise<void> {
    try {
      await this.db
        .update(walEntries)
        .set({
          status,
          updatedAt: Date.now(),
        })
        .where(eq(walEntries.walId, entryId));

      // Check if any rows were affected (SQLite doesn't return rowCount directly)
      const verification = await this.db
        .select({ walId: walEntries.walId })
        .from(walEntries)
        .where(eq(walEntries.walId, entryId))
        .limit(1);

      if (verification.length === 0) {
        throw new Error(`WAL entry ${entryId} not found`);
      }
    } catch (error) {
      throw new Error(
        `Failed to update WAL entry status: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async updateTransactionStatus(
    transactionId: string,
    status: WALTransaction["status"]
  ): Promise<void> {
    try {
      await this.db.run("BEGIN TRANSACTION");

      try {
        // Update transaction status
        await this.db
          .update(walTransactions)
          .set({
            status,
            updatedAt: Date.now(),
          })
          .where(eq(walTransactions.walId, transactionId));

        // Update all related operation statuses
        await this.db
          .update(walEntries)
          .set({
            status,
            updatedAt: Date.now(),
          })
          .where(eq(walEntries.transactionId, transactionId));

        await this.db.run("COMMIT");
      } catch (error) {
        await this.db.run("ROLLBACK");
        throw error;
      }
    } catch (error) {
      throw new Error(
        `Failed to update WAL transaction status: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async fetchUncommittedEntries(
    sessionId?: string
  ): Promise<WALEntry[]> {
    try {
      const whereConditions = [eq(walEntries.status, "pending")];
      if (sessionId) {
        whereConditions.push(eq(walEntries.sessionId, sessionId));
      }

      const query = this.db
        .select()
        .from(walEntries)
        .where(and(...whereConditions))
        .orderBy(walEntries.timestamp)
        .limit(this.config.batchSize);

      const entries = await query;

      return entries.map(this.mapToWALEntry);
    } catch (error) {
      throw new Error(
        `Failed to fetch uncommitted entries: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async fetchEntriesSince(
    timestamp: Date,
    sessionId?: string
  ): Promise<WALEntry[]> {
    try {
      const whereConditions = [gte(walEntries.timestamp, timestamp.getTime())];
      if (sessionId) {
        whereConditions.push(eq(walEntries.sessionId, sessionId));
      }

      const query = this.db
        .select()
        .from(walEntries)
        .where(and(...whereConditions))
        .orderBy(walEntries.timestamp)
        .limit(this.config.batchSize);

      const entries = await query;

      return entries.map(this.mapToWALEntry);
    } catch (error) {
      throw new Error(
        `Failed to fetch entries since timestamp: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async persistRecoveryPoint(
    recoveryPoint: WALRecoveryPoint
  ): Promise<void> {
    try {
      await this.db.insert(walRecoveryPoints).values({
        id: `sqlite_recovery_${Date.now()}_${Math.random().toString(36).substring(2)}`,
        walId: recoveryPoint.id,
        sessionId: recoveryPoint.sessionId,
        walEntryId: recoveryPoint.walEntryId,
        timestamp: recoveryPoint.timestamp.getTime(),
        createdAt: Date.now(),
      });
    } catch (error) {
      throw new Error(
        `Failed to persist recovery point: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async deleteCommittedEntries(olderThan: Date): Promise<number> {
    try {
      await this.db.run("BEGIN TRANSACTION");

      try {
        // First, get the count of entries to be deleted
        const toDelete = await this.db
          .select({ walId: walEntries.walId })
          .from(walEntries)
          .where(
            and(
              eq(walEntries.status, "committed"),
              lt(walEntries.timestamp, olderThan.getTime())
            )
          );

        // Delete entries
        await this.db
          .delete(walEntries)
          .where(
            and(
              eq(walEntries.status, "committed"),
              lt(walEntries.timestamp, olderThan.getTime())
            )
          );

        // Delete related transactions
        await this.db
          .delete(walTransactions)
          .where(
            and(
              eq(walTransactions.status, "committed"),
              lt(walTransactions.timestamp, olderThan.getTime())
            )
          );

        await this.db.run("COMMIT");

        return toDelete.length;
      } catch (error) {
        await this.db.run("ROLLBACK");
        throw error;
      }
    } catch (error) {
      throw new Error(
        `Failed to delete committed entries: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  protected async deleteSessionEntries(sessionId: string): Promise<void> {
    try {
      await this.db.run("BEGIN TRANSACTION");

      try {
        await Promise.all([
          this.db.delete(walEntries).where(eq(walEntries.sessionId, sessionId)),
          this.db
            .delete(walTransactions)
            .where(eq(walTransactions.sessionId, sessionId)),
          this.db
            .delete(walRecoveryPoints)
            .where(eq(walRecoveryPoints.sessionId, sessionId)),
        ]);

        await this.db.run("COMMIT");
      } catch (error) {
        await this.db.run("ROLLBACK");
        throw error;
      }
    } catch (error) {
      throw new Error(
        `Failed to delete session entries: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  private mapToWALEntry(entry: any): WALEntry {
    return {
      id: entry.walId,
      timestamp: new Date(entry.timestamp),
      sessionId: entry.sessionId,
      collection: entry.collection,
      operation: entry.operation as WALEntry["operation"],
      data: JSON.parse(entry.data),
      previousData: entry.previousData
        ? JSON.parse(entry.previousData)
        : undefined,
      transactionId: entry.transactionId || undefined,
      status: entry.status as WALEntry["status"],
      checksum: entry.checksum || undefined,
      metadata: entry.metadata ? JSON.parse(entry.metadata) : undefined,
    };
  }

  /**
   * Get WAL statistics for monitoring
   */
  async getStatistics(): Promise<{
    totalEntries: number;
    pendingEntries: number;
    committedEntries: number;
    rolledBackEntries: number;
    totalTransactions: number;
  }> {
    try {
      const [entryStats, transactionCount] = await Promise.all([
        this.db
          .select({
            status: walEntries.status,
            count: sql`COUNT(*)`,
          })
          .from(walEntries)
          .groupBy(walEntries.status),
        this.db.select({ count: sql`COUNT(*)` }).from(walTransactions),
      ]);

      const stats = {
        totalEntries: 0,
        pendingEntries: 0,
        committedEntries: 0,
        rolledBackEntries: 0,
        totalTransactions: Number(transactionCount[0]?.count) || 0,
      };

      for (const stat of entryStats) {
        const count = stat.count as number;
        stats.totalEntries += count;

        switch (stat.status) {
          case "pending":
            stats.pendingEntries = count;
            break;
          case "committed":
            stats.committedEntries = count;
            break;
          case "rollback":
            stats.rolledBackEntries = count;
            break;
        }
      }

      return stats;
    } catch (error) {
      throw new Error(
        `Failed to get WAL statistics: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }

  /**
   * Perform recovery operations by replaying uncommitted entries
   */
  async performRecovery(
    sessionId: string,
    applyMutation: (entry: WALEntry) => Promise<void>
  ): Promise<number> {
    const uncommittedEntries = await this.getUncommittedEntries(sessionId);
    let appliedCount = 0;

    await this.db.run("BEGIN TRANSACTION");

    try {
      for (const entry of uncommittedEntries) {
        try {
          await applyMutation(entry);
          await this.commitEntry(entry.id);
          appliedCount++;
        } catch (error) {
          console.error(`Failed to recover WAL entry ${entry.id}:`, error);
          await this.rollbackEntry(entry.id);
        }
      }

      await this.db.run("COMMIT");
    } catch (error) {
      await this.db.run("ROLLBACK");
      throw error;
    }

    return appliedCount;
  }

  /**
   * Vacuum WAL tables to reclaim space
   */
  async vacuum(): Promise<void> {
    try {
      await this.db.run("VACUUM");
    } catch (error) {
      throw new Error(
        `Failed to vacuum WAL database: ${error instanceof Error ? error.message : "Unknown error"}`
      );
    }
  }
}