From 88ec15bded6eeb82619cdedb7390f5e7e0bf2386 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Sun, 28 Sep 2025 23:59:38 -0400 Subject: [PATCH] improvement: cache validation requests for sql linter --- frontend/src/__tests__/lru.test.ts | 74 ------------------- .../codemirror/language/languages/sql/sql.ts | 1 + .../src/core/datasets/request-registry.ts | 27 ++++--- .../core/network/CachingRequestRegistry.ts | 74 +++++++++++++++++++ .../core/network/DeferredRequestRegistry.ts | 4 +- .../__tests__/CachingRequestRegistry.test.ts | 73 ++++++++++++++++++ frontend/src/utils/lru.ts | 4 + 7 files changed, 172 insertions(+), 85 deletions(-) delete mode 100644 frontend/src/__tests__/lru.test.ts create mode 100644 frontend/src/core/network/CachingRequestRegistry.ts create mode 100644 frontend/src/core/network/__tests__/CachingRequestRegistry.test.ts diff --git a/frontend/src/__tests__/lru.test.ts b/frontend/src/__tests__/lru.test.ts deleted file mode 100644 index c57341d4b4f..00000000000 --- a/frontend/src/__tests__/lru.test.ts +++ /dev/null @@ -1,74 +0,0 @@ -/* Copyright 2024 Marimo. All rights reserved. */ -import { describe, expect, test } from "vitest"; -import { LRUCache } from "../utils/lru"; - -describe("LRUCache", () => { - test("basic get/set operations", () => { - const cache = new LRUCache(3); - cache.set("a", 1); - cache.set("b", 2); - cache.set("c", 3); - - expect(cache.get("a")).toBe(1); - expect(cache.get("b")).toBe(2); - expect(cache.get("c")).toBe(3); - }); - - test("evicts least recently used item when over capacity", () => { - const cache = new LRUCache(3); - cache.set("a", 1); - cache.set("b", 2); - cache.set("c", 3); - cache.set("d", 4); - - // "a" should be evicted as it was the least recently used - expect(cache.get("a")).toBeUndefined(); - expect(cache.get("b")).toBe(2); - expect(cache.get("c")).toBe(3); - expect(cache.get("d")).toBe(4); - }); - - test("accessing an item makes it most recently used", () => { - const cache = new LRUCache(3); - cache.set("a", 1); - cache.set("b", 2); - cache.set("c", 3); - - // Access "a", making it most recently used - cache.get("a"); - cache.set("d", 4); - - // "b" should be evicted as it becomes the least recently used - expect(cache.get("b")).toBeUndefined(); - expect(cache.get("a")).toBe(1); - expect(cache.get("c")).toBe(3); - expect(cache.get("d")).toBe(4); - }); - - test("updating existing key maintains order", () => { - const cache = new LRUCache(3); - cache.set("a", 1); - cache.set("b", 2); - cache.set("c", 3); - - // Update "a" - cache.set("a", 10); - cache.set("d", 4); - - // "b" should be evicted as "a" was moved to most recent - expect(cache.get("b")).toBeUndefined(); - expect(cache.get("a")).toBe(10); - expect(cache.get("c")).toBe(3); - expect(cache.get("d")).toBe(4); - }); - - test("keys() returns iterator in order", () => { - const cache = new LRUCache(3); - cache.set("a", 1); - cache.set("b", 2); - cache.set("c", 3); - - const keys = [...cache.keys()]; - expect(keys).toEqual(["a", "b", "c"]); - }); -}); diff --git a/frontend/src/core/codemirror/language/languages/sql/sql.ts b/frontend/src/core/codemirror/language/languages/sql/sql.ts index 9d47d1b4847..065405b0c1a 100644 --- a/frontend/src/core/codemirror/language/languages/sql/sql.ts +++ b/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, }); diff --git a/frontend/src/core/datasets/request-registry.ts b/frontend/src/core/datasets/request-registry.ts index 4b1e324703b..18c173586ac 100644 --- a/frontend/src/core/datasets/request-registry.ts +++ b/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, - ValidateSQLResult ->("validate-sql", async (requestId, req) => { - const client = getRequestClient(); - await client.validateSQL({ - requestId: requestId, - ...req, - }); -}); +export const ValidateSQL = new CachingRequestRegistry( + new DeferredRequestRegistry< + Omit, + ValidateSQLResult + >("validate-sql", async (requestId, req) => { + const client = getRequestClient(); + await client.validateSQL({ + requestId: requestId, + ...req, + }); + }), + { + // Only keep the last 3 validation results + maxSize: 3, + }, +); diff --git a/frontend/src/core/network/CachingRequestRegistry.ts b/frontend/src/core/network/CachingRequestRegistry.ts new file mode 100644 index 00000000000..87cce30563e --- /dev/null +++ b/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 = (request: REQ) => string; + +interface CachingOptions { + toKey?: ToKey; + 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 { + private delegate: DeferredRequestRegistry; + private toKey: ToKey; + private cache: LRUCache>; + + static jsonStringifySortKeys(): ToKey { + return (o: T) => { + if (typeof o !== "object" || o === null) { + return String(o); + } + return JSON.stringify(o, Object.keys(o).sort(), 2); + }; + } + + constructor( + delegate: DeferredRequestRegistry, + options: CachingOptions = {}, + ) { + this.delegate = delegate; + this.toKey = + options.toKey ?? CachingRequestRegistry.jsonStringifySortKeys(); + const maxSize = options.maxSize ?? 128; + this.cache = new LRUCache>(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 { + 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); + } +} diff --git a/frontend/src/core/network/DeferredRequestRegistry.ts b/frontend/src/core/network/DeferredRequestRegistry.ts index a83b9888d70..40378b4f5f6 100644 --- a/frontend/src/core/network/DeferredRequestRegistry.ts +++ b/frontend/src/core/network/DeferredRequestRegistry.ts @@ -45,7 +45,9 @@ export class DeferredRequestRegistry { async request(opts: REQ): Promise { 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(); } diff --git a/frontend/src/core/network/__tests__/CachingRequestRegistry.test.ts b/frontend/src/core/network/__tests__/CachingRequestRegistry.test.ts new file mode 100644 index 00000000000..171360486c2 --- /dev/null +++ b/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; + let caching: CachingRequestRegistry; + + 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"); + }); +}); diff --git a/frontend/src/utils/lru.ts b/frontend/src/utils/lru.ts index 9607a0ad3a8..d6537caae46 100644 --- a/frontend/src/utils/lru.ts +++ b/frontend/src/utils/lru.ts @@ -65,4 +65,8 @@ export class LRUCache { public entries() { return this.cache.entries(); } + + public delete(key: K) { + this.cache.delete(key); + } }