improvement: cache validation requests for sql linter

frontend/src/core/codemirror/language/languages/sql/sql.ts

@@ -687,6 +687,7 @@ function sqlValidationExtension(): Extension {
       try {
         const result = await ValidateSQL.request({
           onlyParse: false,
+          dialect: connectionNameToParserDialect(connectionName),
           engine: connectionName,
           query: sqlContent,
         });

frontend/src/core/datasets/request-registry.ts

@@ -4,6 +4,7 @@ import type {
   SQLTablePreview,
   ValidateSQLResult,
 } from "../kernel/messages";
+import { CachingRequestRegistry } from "../network/CachingRequestRegistry";
 import { DeferredRequestRegistry } from "../network/DeferredRequestRegistry";
 import { getRequestClient } from "../network/requests";
 import type {
@@ -38,13 +39,19 @@ export const PreviewSQLTableList = new DeferredRequestRegistry<
   });
 });
 
-export const ValidateSQL = new DeferredRequestRegistry<
-  Omit<ValidateSQLRequest, "requestId">,
-  ValidateSQLResult
->("validate-sql", async (requestId, req) => {
-  const client = getRequestClient();
-  await client.validateSQL({
-    requestId: requestId,
-    ...req,
-  });
-});
+export const ValidateSQL = new CachingRequestRegistry(
+  new DeferredRequestRegistry<
+    Omit<ValidateSQLRequest, "requestId">,
+    ValidateSQLResult
+  >("validate-sql", async (requestId, req) => {
+    const client = getRequestClient();
+    await client.validateSQL({
+      requestId: requestId,
+      ...req,
+    });
+  }),
+  {
+    // Only keep the last 3 validation results
+    maxSize: 3,
+  },
+);

frontend/src/core/network/CachingRequestRegistry.ts

@@ -0,0 +1,74 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { LRUCache } from "@/utils/lru";
+import type {
+  DeferredRequestRegistry,
+  RequestId,
+} from "./DeferredRequestRegistry";
+
+type ToKey<REQ> = (request: REQ) => string;
+
+interface CachingOptions<REQ> {
+  toKey?: ToKey<REQ>;
+  maxSize?: number;
+}
+
+/**
+ * Light wrapper adding memoization and in-flight de-duplication on top of
+ * DeferredRequestRegistry, keyed by a string representation of the request.
+ */
+export class CachingRequestRegistry<REQ, RES> {
+  private delegate: DeferredRequestRegistry<REQ, RES>;
+  private toKey: ToKey<REQ>;
+  private cache: LRUCache<string, Promise<RES>>;
+
+  static jsonStringifySortKeys<T>(): ToKey<T> {
+    return (o: T) => {
+      if (typeof o !== "object" || o === null) {
+        return String(o);
+      }
+      return JSON.stringify(o, Object.keys(o).sort(), 2);
+    };
+  }
+
+  constructor(
+    delegate: DeferredRequestRegistry<REQ, RES>,
+    options: CachingOptions<REQ> = {},
+  ) {
+    this.delegate = delegate;
+    this.toKey =
+      options.toKey ?? CachingRequestRegistry.jsonStringifySortKeys();
+    const maxSize = options.maxSize ?? 128;
+    this.cache = new LRUCache<string, Promise<RES>>(maxSize);
+  }
+
+  /**
+   * Resolve via cache if present, else delegate; de-duplicates concurrent
+   * requests with the same key and stores successful results in the cache.
+   */
+  public request(req: REQ): Promise<RES> {
+    const key = this.toKey(req);
+
+    const cached = this.cache.get(key);
+    if (cached !== undefined) {
+      return cached;
+    }
+
+    const promise = this.delegate.request(req);
+    this.cache.set(key, promise);
+    return promise.catch((err) => {
+      this.cache.delete(key);
+      throw err;
+    });
+  }
+
+  // Path through to the delegate
+  public resolve(requestId: RequestId, response: RES) {
+    this.delegate.resolve(requestId, response);
+  }
+
+  // Path through to the delegate
+  public reject(requestId: RequestId, error: Error) {
+    this.delegate.reject(requestId, error);
+  }
+}

frontend/src/core/network/DeferredRequestRegistry.ts

@@ -45,7 +45,9 @@ export class DeferredRequestRegistry<REQ, RES> {
   async request(opts: REQ): Promise<RES> {
     if (this.opts.resolveExistingRequests) {
       const result = this.opts.resolveExistingRequests();
-      this.requests.forEach((deferred) => deferred.resolve(result));
+      for (const deferred of this.requests.values()) {
+        deferred.resolve(result);
+      }
       this.requests.clear();
     }
 

frontend/src/core/network/__tests__/CachingRequestRegistry.test.ts

@@ -0,0 +1,73 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { CachingRequestRegistry } from "../CachingRequestRegistry";
+import {
+  DeferredRequestRegistry,
+  type RequestId,
+} from "../DeferredRequestRegistry";
+
+vi.mock("@/utils/uuid", () => ({
+  generateUUID: vi.fn().mockReturnValue("uuid"),
+}));
+
+describe("CachingRequestRegistry", () => {
+  const REQUEST_ID = "uuid" as RequestId;
+  let makeRequestMock = vi.fn();
+  let delegate: DeferredRequestRegistry<unknown, unknown>;
+  let caching: CachingRequestRegistry<unknown, unknown>;
+
+  beforeEach(() => {
+    makeRequestMock = vi.fn().mockResolvedValue(undefined);
+    delegate = new DeferredRequestRegistry("operation", makeRequestMock);
+    caching = new CachingRequestRegistry(delegate);
+  });
+
+  it("should cache successful responses for identical requests", async () => {
+    const req = { a: 1 };
+
+    const p1 = caching.request(req);
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+    expect(makeRequestMock).toHaveBeenCalledWith(REQUEST_ID, req);
+
+    // Resolve first request
+    delegate.resolve(REQUEST_ID, "response");
+    await expect(p1).resolves.toBe("response");
+
+    // Second call with equivalent request gets served from cache
+    const p2 = caching.request({ a: 1 });
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+    await expect(p2).resolves.toBe("response");
+  });
+
+  it("should de-duplicate in-flight requests with the same key", async () => {
+    const req = { q: "select *" };
+
+    const p1 = caching.request(req);
+    const p2 = caching.request({ q: "select *" });
+
+    // Only one network invocation while in-flight
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+    expect(p1).toStrictEqual(p2);
+
+    // Resolve and ensure both resolve to same result
+    delegate.resolve(REQUEST_ID, "ok");
+    await expect(p1).resolves.toBe("ok");
+    await expect(p2).resolves.toBe("ok");
+  });
+
+  it("should not cache errors", async () => {
+    // First call rejects
+    makeRequestMock.mockRejectedValueOnce(new Error("boom"));
+
+    await expect(caching.request({ x: 1 })).rejects.toThrow("boom");
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+
+    // Next call should attempt again (not cached)
+    const p2 = caching.request({ x: 1 });
+    expect(makeRequestMock).toHaveBeenCalledTimes(2);
+
+    // Resolve the second request
+    delegate.resolve(REQUEST_ID, "ok");
+    await expect(p2).resolves.toBe("ok");
+  });
+});

frontend/src/utils/lru.ts

@@ -65,4 +65,8 @@ export class LRUCache<K, V> {
   public entries() {
     return this.cache.entries();
   }
+
+  public delete(key: K) {
+    this.cache.delete(key);
+  }
 }