dzql 0.1.6 → 0.2.0

package/src/server/index.js

@@ -2,12 +2,96 @@ import { createWebSocketHandlers, verify_jwt_token } from "./ws.js";
 import { closeConnections, setupListeners, sql, db } from "./db.js";
 import * as defaultApi from "./api.js";
 import { serverLogger, notifyLogger } from "./logger.js";
+import { getSubscriptionsBySubscribable, paramsMatch } from "./subscriptions.js";
 
 // Re-export commonly used utilities
 export { sql, db } from "./db.js";
 export { metaRoute } from "./meta-route.js";
 export { createMCPRoute } from "./mcp.js";
 
+/**
+ * Process subscription updates when a database event occurs
+ * Checks if any active subscriptions are affected and sends updates
+ * @param {Object} event - Database event {table, op, pk, before, after}
+ * @param {Function} broadcast - Broadcast function from WebSocket handlers
+ */
+async function processSubscriptionUpdates(event, broadcast) {
+  const { table, op, before, after } = event;
+
+  // Get all active subscriptions grouped by subscribable
+  const subscriptionsByName = getSubscriptionsBySubscribable();
+
+  if (subscriptionsByName.size === 0) {
+    return; // No active subscriptions
+  }
+
+  notifyLogger.debug(`Checking ${subscriptionsByName.size} subscribable(s) for affected subscriptions`);
+
+  // For each unique subscribable, check if this event affects any subscriptions
+  for (const [subscribableName, subs] of subscriptionsByName.entries()) {
+    try {
+      // Ask PostgreSQL which subscription instances are affected
+      const result = await db.query(
+        `SELECT ${subscribableName}_affected_documents($1, $2, $3, $4) as affected`,
+        [table, op, before, after]
+      );
+
+      const affectedParamSets = result.rows[0]?.affected;
+
+      if (!affectedParamSets || affectedParamSets.length === 0) {
+        continue; // This subscribable not affected
+      }
+
+      notifyLogger.debug(`${subscribableName}: ${affectedParamSets.length} param set(s) affected`);
+
+      // Match affected params to active subscriptions
+      for (const affectedParams of affectedParamSets) {
+        for (const sub of subs) {
+          // Check if this subscription matches the affected params
+          if (paramsMatch(sub.params, affectedParams)) {
+            try {
+              // Re-execute query to get updated data
+              const updated = await db.query(
+                `SELECT get_${subscribableName}($1, $2) as data`,
+                [sub.params, sub.user_id]
+              );
+
+              const data = updated.rows[0]?.data;
+
+              // Send update to specific connection
+              const message = JSON.stringify({
+                jsonrpc: "2.0",
+                method: "subscription:update",
+                params: {
+                  subscription_id: sub.subscriptionId,
+                  subscribable: subscribableName,
+                  data
+                }
+              });
+
+              const sent = broadcast.toConnection(sub.connection_id, message);
+              if (sent) {
+                notifyLogger.debug(`Sent update to subscription ${sub.subscriptionId.slice(0, 8)}...`);
+              } else {
+                notifyLogger.warn(`Failed to send update to connection ${sub.connection_id.slice(0, 8)}...`);
+              }
+            } catch (error) {
+              notifyLogger.error(`Failed to update subscription ${sub.subscriptionId}:`, error.message);
+            }
+          }
+        }
+      }
+    } catch (error) {
+      // If the subscribable function doesn't exist, just skip
+      if (error.message && error.message.includes('does not exist')) {
+        notifyLogger.debug(`Subscribable ${subscribableName} functions not found, skipping`);
+      } else {
+        notifyLogger.error(`Error processing subscriptions for ${subscribableName}:`, error.message);
+      }
+    }
+  }
+}
+
 /**
  * Create a DZQL server with WebSocket support, real-time updates, and automatic CRUD operations
  *
@@ -97,10 +181,11 @@ export function createServer(options = {}) {
   });
 
   // Setup NOTIFY listeners for real-time events
-  setupListeners((event) => {
+  setupListeners(async (event) => {
     // Handle single dzql event with filtering
     const { notify_users, ...eventData } = event;
 
+    // PATTERN 2: Need to Know notifications (existing)
     // Create JSON-RPC notification
     const message = JSON.stringify({
       jsonrpc: "2.0",
@@ -118,6 +203,10 @@ export function createServer(options = {}) {
       notifyLogger.debug(`Broadcasting ${event.table}:${event.op} to all users`);
       broadcast(message);
     }
+
+    // PATTERN 1: Live Query subscriptions (new)
+    // Check if any subscriptions are affected by this event
+    await processSubscriptionUpdates(event, broadcast);
   });
 
   routes['/health'] = () => new Response("OK", { status: 200 });

package/src/server/subscriptions.js

@@ -0,0 +1,209 @@
+/**
+ * Subscription Manager for Live Query Subscriptions
+ * Manages in-memory subscription registry and matching
+ */
+
+import crypto from 'crypto';
+import { wsLogger } from './logger.js';
+
+/**
+ * In-memory subscription registry
+ * Structure: subscription_id -> { subscribable, user_id, connection_id, params }
+ */
+const subscriptions = new Map();
+
+/**
+ * Track subscriptions by connection for cleanup
+ * Structure: connection_id -> Set<subscription_id>
+ */
+const connectionSubscriptions = new Map();
+
+/**
+ * Register a new subscription
+ * @param {string} subscribableName - Name of the subscribable
+ * @param {number} userId - User ID
+ * @param {string} connectionId - WebSocket connection ID
+ * @param {object} params - Subscription parameters
+ * @returns {string} - Subscription ID
+ */
+export function registerSubscription(subscribableName, userId, connectionId, params) {
+  const subscriptionId = crypto.randomUUID();
+
+  // Store subscription
+  subscriptions.set(subscriptionId, {
+    subscribable: subscribableName,
+    user_id: userId,
+    connection_id: connectionId,
+    params,
+    created_at: new Date()
+  });
+
+  // Track by connection
+  if (!connectionSubscriptions.has(connectionId)) {
+    connectionSubscriptions.set(connectionId, new Set());
+  }
+  connectionSubscriptions.get(connectionId).add(subscriptionId);
+
+  wsLogger.debug(`Subscription registered: ${subscriptionId.slice(0, 8)}... (${subscribableName})`, {
+    user_id: userId,
+    params
+  });
+
+  return subscriptionId;
+}
+
+/**
+ * Unregister a subscription
+ * @param {string} subscriptionId - Subscription ID to remove
+ * @returns {boolean} - True if subscription was found and removed
+ */
+export function unregisterSubscription(subscriptionId) {
+  const sub = subscriptions.get(subscriptionId);
+
+  if (!sub) {
+    wsLogger.debug(`Subscription not found: ${subscriptionId}`);
+    return false;
+  }
+
+  // Remove from connection tracking
+  const connSubs = connectionSubscriptions.get(sub.connection_id);
+  if (connSubs) {
+    connSubs.delete(subscriptionId);
+    if (connSubs.size === 0) {
+      connectionSubscriptions.delete(sub.connection_id);
+    }
+  }
+
+  // Remove subscription
+  subscriptions.delete(subscriptionId);
+
+  wsLogger.debug(`Subscription removed: ${subscriptionId.slice(0, 8)}...`);
+  return true;
+}
+
+/**
+ * Unregister subscription by params
+ * Useful for unsubscribe_* methods that specify params instead of subscription ID
+ * @param {string} subscribableName - Name of the subscribable
+ * @param {string} connectionId - WebSocket connection ID
+ * @param {object} params - Subscription parameters to match
+ * @returns {boolean} - True if subscription was found and removed
+ */
+export function unregisterSubscriptionByParams(subscribableName, connectionId, params) {
+  const paramsStr = JSON.stringify(params);
+
+  // Find matching subscription
+  for (const [subId, sub] of subscriptions.entries()) {
+    if (
+      sub.subscribable === subscribableName &&
+      sub.connection_id === connectionId &&
+      JSON.stringify(sub.params) === paramsStr
+    ) {
+      return unregisterSubscription(subId);
+    }
+  }
+
+  wsLogger.debug(`No matching subscription found for ${subscribableName} with params`, params);
+  return false;
+}
+
+/**
+ * Remove all subscriptions for a connection
+ * Called when WebSocket connection closes
+ * @param {string} connectionId - Connection ID
+ * @returns {number} - Number of subscriptions removed
+ */
+export function removeConnectionSubscriptions(connectionId) {
+  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 for a specific subscribable
+ * @param {string} subscribableName - Name of the subscribable
+ * @returns {Array} - Array of {subscriptionId, subscription} objects
+ */
+export function getSubscriptionsByName(subscribableName) {
+  const result = [];
+
+  for (const [subId, sub] of subscriptions.entries()) {
+    if (sub.subscribable === subscribableName) {
+      result.push({ subscriptionId: subId, ...sub });
+    }
+  }
+
+  return result;
+}
+
+/**
+ * Get all subscriptions grouped by subscribable name
+ * @returns {Map<string, Array>} - Map of subscribable name to array of subscriptions
+ */
+export function getSubscriptionsBySubscribable() {
+  const grouped = new Map();
+
+  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)
+ * @param {object} a - First params object
+ * @param {object} b - Second params object
+ * @returns {boolean} - True if params match
+ */
+export function paramsMatch(a, b) {
+  return JSON.stringify(a) === JSON.stringify(b);
+}
+
+/**
+ * Get subscription statistics
+ * @returns {object} - Stats object
+ */
+export function getStats() {
+  return {
+    total_subscriptions: subscriptions.size,
+    active_connections: connectionSubscriptions.size,
+    subscriptions_by_subscribable: Array.from(getSubscriptionsBySubscribable().entries()).map(
+      ([name, subs]) => ({
+        name,
+        count: subs.length
+      })
+    )
+  };
+}
+
+/**
+ * Get all active subscriptions (for debugging)
+ * @returns {Array} - Array of all subscriptions
+ */
+export function getAllSubscriptions() {
+  return Array.from(subscriptions.entries()).map(([id, sub]) => ({
+    id,
+    ...sub
+  }));
+}

package/src/server/ws.js

@@ -6,6 +6,12 @@ import {
   db,
 } from "./db.js";
 import { wsLogger, authLogger } from "./logger.js";
+import {
+  registerSubscription,
+  unregisterSubscription,
+  unregisterSubscriptionByParams,
+  removeConnectionSubscriptions
+} from "./subscriptions.js";
 
 // Environment configuration
 const JWT_SECRET_STRING = process.env.JWT_SECRET;
@@ -304,6 +310,57 @@ export function createRPCHandler(customHandlers = {}) {
         return create_rpc_response(id, result);
       }
 
+      // SUBSCRIPTION HANDLERS - Pattern match on method name
+      if (method.startsWith("subscribe_")) {
+        const subscribableName = method.replace("subscribe_", "");
+        wsLogger.debug(`Subscribe request: ${subscribableName}`, params);
+
+        try {
+          // Execute initial query (this also checks permissions)
+          const queryResult = await db.query(
+            `SELECT get_${subscribableName}($1, $2) as data`,
+            [params, ws.data.user_id]
+          );
+
+          const data = queryResult.rows[0]?.data;
+
+          // Register subscription in memory
+          const subscriptionId = registerSubscription(
+            subscribableName,
+            ws.data.user_id,
+            ws.data.connection_id,
+            params
+          );
+
+          const result = {
+            subscription_id: subscriptionId,
+            data
+          };
+
+          wsLogger.response(method, result, Date.now() - startTime);
+          return create_rpc_response(id, result);
+        } catch (error) {
+          wsLogger.error(`Subscribe failed for ${subscribableName}:`, error.message);
+          return create_rpc_error(id, -32603, error.message);
+        }
+      }
+
+      if (method.startsWith("unsubscribe_")) {
+        const subscribableName = method.replace("unsubscribe_", "");
+        wsLogger.debug(`Unsubscribe request: ${subscribableName}`, params);
+
+        // Remove subscription by params
+        const removed = unregisterSubscriptionByParams(
+          subscribableName,
+          ws.data.connection_id,
+          params
+        );
+
+        const result = { success: removed };
+        wsLogger.response(method, result, Date.now() - startTime);
+        return create_rpc_response(id, result);
+      }
+
       // Check for custom handlers
       if (customHandlers[method]) {
         wsLogger.debug(`Calling custom handler: ${method}`);
@@ -417,7 +474,14 @@ export function createWebSocketHandlers(options = {}) {
     close(ws) {
       const id = ws.data.connection_id;
       connections.delete(id);
-      wsLogger.info(`Connection closed: ${id?.slice(0, 8)}...`);
+
+      // Clean up all subscriptions for this connection
+      const removedCount = removeConnectionSubscriptions(id);
+      if (removedCount > 0) {
+        wsLogger.info(`Connection closed: ${id?.slice(0, 8)}... (${removedCount} subscriptions removed)`);
+      } else {
+        wsLogger.info(`Connection closed: ${id?.slice(0, 8)}...`);
+      }
 
       // Call custom disconnection handler
       if (onDisconnection) {
@@ -434,7 +498,7 @@ export function createWebSocketHandlers(options = {}) {
 
 // Broadcast message to all authenticated connections or specific client_ids
 export function createBroadcaster(connections) {
-  return function broadcastToConnections(message, client_ids = null) {
+  const broadcastToConnections = function(message, client_ids = null) {
     if (client_ids && Array.isArray(client_ids)) {
       // Send to specific user_ids
       for (const [id, ws] of connections) {
@@ -451,6 +515,18 @@ export function createBroadcaster(connections) {
       }
     }
   };
+
+  // Add helper function to send to a specific connection
+  broadcastToConnections.toConnection = function(connectionId, message) {
+    const ws = connections.get(connectionId);
+    if (ws && ws.readyState === 1) { // 1 = OPEN
+      ws.send(message);
+      return true;
+    }
+    return false;
+  };
+
+  return broadcastToConnections;
 }
 
 // Legacy export for backward compatibility