dzql 0.6.13 → 0.6.15

package/src/runtime/subscriptions.ts

@@ -0,0 +1,189 @@
+/**
+ * Subscription Manager for Live Query Subscriptions
+ * Manages in-memory subscription registry and matching
+ */
+
+import crypto from 'crypto';
+import { wsLogger } from './logger.js';
+
+/** Subscription entry */
+interface Subscription {
+  subscribable: string;
+  user_id: number | null;
+  connection_id: string;
+  params: Record<string, unknown>;
+  created_at: Date;
+}
+
+/** Subscribable metadata */
+interface SubscribableMetadata {
+  scopeTables: string[];
+  pathMapping: Record<string, string>;
+  rootEntity: string | null;
+}
+
+/** In-memory subscription registry: subscription_id -> subscription */
+const subscriptions = new Map<string, Subscription>();
+
+/** Track subscriptions by connection for cleanup: connection_id -> Set<subscription_id> */
+const connectionSubscriptions = new Map<string, Set<string>>();
+
+/** Cache for subscribable metadata: subscribable_name -> metadata */
+const subscribableMetadataCache = new Map<string, SubscribableMetadata>();
+
+/**
+ * Register a new subscription
+ */
+export function registerSubscription(
+  subscribableName: string,
+  userId: number | null,
+  connectionId: string,
+  params: Record<string, unknown>
+): string {
+  const subscriptionId = crypto.randomUUID();
+
+  subscriptions.set(subscriptionId, {
+    subscribable: subscribableName,
+    user_id: userId,
+    connection_id: connectionId,
+    params,
+    created_at: new Date()
+  });
+
+  if (!connectionSubscriptions.has(connectionId)) {
+    connectionSubscriptions.set(connectionId, new Set());
+  }
+  connectionSubscriptions.get(connectionId)!.add(subscriptionId);
+
+  wsLogger.debug(`Subscription registered: ${subscriptionId.slice(0, 8)}... (${subscribableName})`);
+
+  return subscriptionId;
+}
+
+/**
+ * Unregister a subscription
+ */
+export function unregisterSubscription(subscriptionId: string): boolean {
+  const sub = subscriptions.get(subscriptionId);
+
+  if (!sub) {
+    return false;
+  }
+
+  const connSubs = connectionSubscriptions.get(sub.connection_id);
+  if (connSubs) {
+    connSubs.delete(subscriptionId);
+    if (connSubs.size === 0) {
+      connectionSubscriptions.delete(sub.connection_id);
+    }
+  }
+
+  subscriptions.delete(subscriptionId);
+  wsLogger.debug(`Subscription removed: ${subscriptionId.slice(0, 8)}...`);
+  return true;
+}
+
+/**
+ * Unregister subscription by params
+ */
+export function unregisterSubscriptionByParams(
+  subscribableName: string,
+  connectionId: string,
+  params: Record<string, unknown>
+): boolean {
+  const paramsStr = JSON.stringify(params);
+
+  for (const [subId, sub] of subscriptions.entries()) {
+    if (
+      sub.subscribable === subscribableName &&
+      sub.connection_id === connectionId &&
+      JSON.stringify(sub.params) === paramsStr
+    ) {
+      return unregisterSubscription(subId);
+    }
+  }
+
+  return false;
+}
+
+/**
+ * Remove all subscriptions for a connection (called on WebSocket close)
+ */
+export function removeConnectionSubscriptions(connectionId: string): number {
+  const subIds = connectionSubscriptions.get(connectionId);
+
+  if (!subIds) {
+    return 0;
+  }
+
+  let count = 0;
+  for (const subId of subIds) {
+    if (subscriptions.delete(subId)) {
+      count++;
+    }
+  }
+
+  connectionSubscriptions.delete(connectionId);
+
+  if (count > 0) {
+    wsLogger.info(`Removed ${count} subscription(s) for connection ${connectionId.slice(0, 8)}...`);
+  }
+
+  return count;
+}
+
+/**
+ * Get all subscriptions grouped by subscribable name
+ */
+export function getSubscriptionsBySubscribable(): Map<string, Array<Subscription & { subscriptionId: string }>> {
+  const grouped = new Map<string, Array<Subscription & { subscriptionId: string }>>();
+
+  for (const [subId, sub] of subscriptions.entries()) {
+    if (!grouped.has(sub.subscribable)) {
+      grouped.set(sub.subscribable, []);
+    }
+    grouped.get(sub.subscribable)!.push({ subscriptionId: subId, ...sub });
+  }
+
+  return grouped;
+}
+
+/**
+ * Check if params match (deep equality)
+ */
+export function paramsMatch(a: Record<string, unknown>, b: Record<string, unknown>): boolean {
+  return JSON.stringify(a) === JSON.stringify(b);
+}
+
+/**
+ * Cache subscribable metadata (from manifest)
+ */
+export function cacheSubscribableMetadata(subscribableName: string, metadata: SubscribableMetadata): void {
+  subscribableMetadataCache.set(subscribableName, metadata);
+}
+
+/**
+ * Get subscribable metadata from cache
+ */
+export function getSubscribableMetadata(subscribableName: string): SubscribableMetadata | undefined {
+  return subscribableMetadataCache.get(subscribableName);
+}
+
+/**
+ * Get scope tables for a subscribable
+ */
+export function getSubscribableScopeTables(subscribableName: string): string[] {
+  const metadata = subscribableMetadataCache.get(subscribableName);
+  return metadata?.scopeTables || [];
+}
+
+/**
+ * Clear metadata cache
+ */
+export function clearSubscribableMetadataCache(subscribableName?: string): void {
+  if (subscribableName) {
+    subscribableMetadataCache.delete(subscribableName);
+  } else {
+    subscribableMetadataCache.clear();
+  }
+}

package/src/runtime/ws.ts

@@ -1,18 +1,23 @@
 import { ServerWebSocket } from "bun";
-import { handleRequest } from "./server.js"; // The secure router
+import { handleRequest } from "./server.js";
 import { verifyToken, signToken } from "./auth.js";
 import { Database } from "./db.js";
 import { wsLogger, authLogger, logWsAccess } from "./logger.js";
+import { getManifest } from "./manifest_loader.js";
+import {
+  registerSubscription,
+  removeConnectionSubscriptions,
+  unregisterSubscriptionByParams,
+  cacheSubscribableMetadata
+} from "./subscriptions.js";
 
-// WebSocket configuration
-const WS_MAX_MESSAGE_SIZE = parseInt(process.env.WS_MAX_MESSAGE_SIZE || "1048576", 10); // 1MB default
+const WS_MAX_MESSAGE_SIZE = parseInt(process.env.WS_MAX_MESSAGE_SIZE || "1048576", 10);
 
 interface WSContext {
-  id?: string;  // Set in handleOpen
+  id?: string;
   userId?: number;
-  subscriptions?: Set<string>; // Set of subscription IDs, set in handleOpen
-  lastPing?: number;  // Set in handleOpen
-  token?: string; // Token passed from URL query params during upgrade
+  lastPing?: number;
+  token?: string;
 }
 
 export class WebSocketServer {
@@ -23,23 +28,18 @@ export class WebSocketServer {
     this.db = db;
   }
 
-  // Bun.serve websocket configuration object
   get websocket() {
     return {
-      // WebSocket options for Bun
       perMessageDeflate: true,
       maxPayloadLength: WS_MAX_MESSAGE_SIZE,
-      idleTimeout: 0, // No idle timeout - realtime connections stay open indefinitely
-
-      // Handler hooks
+      idleTimeout: 0,
       open: (ws: ServerWebSocket<WSContext>) => this.handleOpen(ws),
       message: (ws: ServerWebSocket<WSContext>, message: string) => this.handleMessage(ws, message),
       close: (ws: ServerWebSocket<WSContext>) => this.handleClose(ws),
-      drain: (ws: ServerWebSocket<WSContext>) => {}
+      drain: () => {}
     };
   }
 
-  // Legacy alias for backwards compatibility
   get handlers() {
     return this.websocket;
   }
@@ -50,16 +50,11 @@ export class WebSocketServer {
 
     ws.data = {
       id,
-      subscriptions: new Set(),
       lastPing: Date.now()
     };
     this.connections.set(id, ws);
     wsLogger.info(`Client ${id} connected`);
 
-    // Subscribe to global broadcast channel initially
-    ws.subscribe("broadcast");
-
-    // If token was provided in URL, verify and authenticate
     let user: any = null;
     if (token) {
       try {
@@ -67,26 +62,21 @@ export class WebSocketServer {
         ws.data.userId = session.userId;
         authLogger.info(`Client ${id} authenticated via token as user ${session.userId}`);
 
-        // Fetch user profile using get_users
         try {
           const profile = await handleRequest(this.db, 'get_users', { id: session.userId }, session.userId);
           if (profile) {
-            // Remove sensitive data
             const { password_hash, ...safeProfile } = profile;
             user = safeProfile;
           }
         } catch (e: any) {
           authLogger.error(`Failed to fetch profile for user ${session.userId}:`, e.message);
-          // Still authenticated, just no profile
           user = { id: session.userId };
         }
       } catch (e: any) {
         authLogger.debug(`Client ${id} token verification failed:`, e.message);
-        // Token invalid, user remains null (anonymous connection)
       }
     }
 
-    // Send connection:ready message
     ws.send(JSON.stringify({
       jsonrpc: "2.0",
       method: "connection:ready",
@@ -95,7 +85,7 @@ export class WebSocketServer {
   }
 
   private async handleMessage(ws: ServerWebSocket<WSContext>, message: string) {
-    ws.data.lastPing = Date.now(); // Update activity
+    ws.data.lastPing = Date.now();
     const start = Date.now();
     let method = "unknown";
 
@@ -103,14 +93,11 @@ export class WebSocketServer {
       const req = JSON.parse(message);
       method = req.method || "unknown";
 
-      // Handle Ping (Client-side heartbeat)
       if (req.method === 'ping') {
         ws.send(JSON.stringify({ id: req.id, result: 'pong' }));
-        wsLogger.trace(`Client ${ws.data.id} ping`);
         return;
       }
 
-      // Handle Auth Handshake
       if (req.method === 'auth') {
         try {
           const session = await verifyToken(req.params.token);
@@ -125,31 +112,40 @@ export class WebSocketServer {
         return;
       }
 
-      // Require Auth for other methods?
-      // V2 spec says `p_user_id` is passed to functions.
-      // If not auth'd, we can pass null or throw.
       const userId = ws.data.userId || null;
+      const connectionId = ws.data.id!;
 
-      // Handle RPC
       if (req.method) {
         if (req.method.startsWith("subscribe_")) {
-          // Handle Subscription Registration
           const subscribableName = req.method.replace("subscribe_", "");
           const getFnName = `get_${subscribableName}`;
 
           try {
-            // Call the get_ function to fetch initial data
             const snapshot = await handleRequest(this.db, getFnName, req.params, userId);
 
-            // Register subscription
-            const subId = `${subscribableName}:${JSON.stringify(req.params)}`;
-            ws.data.subscriptions?.add(subId);
+            // Get metadata from manifest and cache it
+            const manifest = getManifest();
+            const subscribable = manifest.subscribables[subscribableName];
+            if (subscribable) {
+              cacheSubscribableMetadata(subscribableName, {
+                scopeTables: subscribable.scopeTables || [],
+                pathMapping: {}, // TODO: build from includes
+                rootEntity: subscribable.root?.entity || null
+              });
+            }
+
+            // Register subscription in central registry
+            const subscriptionId = registerSubscription(
+              subscribableName,
+              userId,
+              connectionId,
+              req.params || {}
+            );
 
-            // Return snapshot with subscription_id and schema
             ws.send(JSON.stringify({
               id: req.id,
               result: {
-                subscription_id: subId,
+                subscription_id: subscriptionId,
                 data: snapshot.data,
                 schema: snapshot.schema
               }
@@ -163,17 +159,18 @@ export class WebSocketServer {
             }));
             logWsAccess("SUB", req.method, Date.now() - start, userId, e.message);
           }
+        } else if (req.method.startsWith("unsubscribe_")) {
+          const subscribableName = req.method.replace("unsubscribe_", "");
+          const removed = unregisterSubscriptionByParams(subscribableName, connectionId, req.params || {});
+          ws.send(JSON.stringify({ id: req.id, result: { success: removed } }));
+          logWsAccess("UNSUB", req.method, Date.now() - start, userId);
         } else {
-          // Normal Function Call
           try {
             const result = await handleRequest(this.db, req.method, req.params, userId);
 
-            // Auto-generate token for auth methods
             if (req.method === 'login_user' || req.method === 'register_user') {
               const token = await signToken({ user_id: result.user_id, role: 'user' });
-              // Update connection's userId for subsequent calls
               ws.data.userId = result.user_id;
-              // Return profile + token
               ws.send(JSON.stringify({ id: req.id, result: { ...result, token } }));
               logWsAccess("CALL", req.method, Date.now() - start, result.user_id);
             } else {
@@ -193,7 +190,6 @@ export class WebSocketServer {
     } catch (e: any) {
       wsLogger.error(`Error processing message from ${ws.data.id}:`, e.message);
       logWsAccess("CALL", method, Date.now() - start, ws.data.userId, e.message);
-      // Try to send error back if JSON parse didn't fail
       try {
         const reqId = JSON.parse(message).id;
         ws.send(JSON.stringify({ id: reqId, error: { code: 'INTERNAL_ERROR', message: e.message } }));
@@ -202,20 +198,43 @@ export class WebSocketServer {
   }
 
   private handleClose(ws: ServerWebSocket<WSContext>) {
-    if (ws.data.id) {
-      this.connections.delete(ws.data.id);
+    const connectionId = ws.data.id;
+    if (connectionId) {
+      this.connections.delete(connectionId);
+      removeConnectionSubscriptions(connectionId);
+    }
+    wsLogger.info(`Client ${connectionId} disconnected`);
+  }
+
+  /**
+   * Send message to a specific connection
+   */
+  public toConnection(connectionId: string, message: string): boolean {
+    const ws = this.connections.get(connectionId);
+    if (ws && ws.readyState === 1) { // WebSocket.OPEN
+      ws.send(message);
+      return true;
+    }
+    return false;
+  }
+
+  /**
+   * Get all connection IDs for a specific user
+   */
+  public getConnectionsByUserId(userId: number): string[] {
+    const connectionIds: string[] = [];
+    for (const [id, ws] of this.connections.entries()) {
+      if (ws.data?.userId === userId) {
+        connectionIds.push(id);
+      }
     }
-    wsLogger.info(`Client ${ws.data.id} disconnected`);
+    return connectionIds;
   }
 
-  public broadcast(message: string) {
-    // Use Bun's native publish for efficiency
-    // 'broadcast' topic is subscribed by all on connect
-    // In V2, we might have specific topics per subscription key
-    // server.publish("broadcast", message);
-    // Since this class doesn't hold the 'server' instance directly,
-    // we iterate or need to pass server in.
-    // For now, iteration is fine for prototype.
+  /**
+   * Broadcast to all connections (for backwards compatibility)
+   */
+  public broadcast(message: string): void {
     for (const ws of this.connections.values()) {
       ws.send(message);
     }