@lightdash/warehouses 0.3453.0 → 0.3454.0

package/dist/warehouseClients/DuckdbWarehouseClient.test.js

@@ -91,6 +91,7 @@ const createMockConnection = (streamMock, runMock = vi.fn(), opts) => ({
         run: runMock,
         stream: streamMock,
         extractStatements: opts?.extractStatements ?? createMockExtractStatements(),
+        interrupt: opts?.interrupt ?? vi.fn(),
         closeSync: vi.fn(),
         disconnectSync: vi.fn(),
     }),
@@ -340,7 +341,99 @@ describe('DuckdbWarehouseClient', () => {
         expect(runMock).not.toHaveBeenCalledWith("SET memory_limit = '6GB';");
         expect(runMock).not.toHaveBeenCalledWith('SET threads = 8;');
     });
-    it('should log structured DuckDB profile metrics with query tags', async () => {
+    it('hardens each query connection on a shared instance', async () => {
+        const bootstrapRunMock = vi.fn();
+        const queryRunMock = vi.fn();
+        const streamMock = vi.fn(async () => getMockStreamResult([[{ val: 1 }]], [DUCKDB_TYPE_IDS.INTEGER]));
+        const connectMock = vi
+            .fn()
+            .mockResolvedValueOnce({
+            run: bootstrapRunMock,
+            closeSync: vi.fn(),
+            disconnectSync: vi.fn(),
+        })
+            .mockResolvedValueOnce({
+            run: queryRunMock,
+            stream: streamMock,
+            extractStatements: createMockExtractStatements(),
+            closeSync: vi.fn(),
+            disconnectSync: vi.fn(),
+        });
+        createInstanceMock.mockResolvedValue({
+            connect: connectMock,
+            closeSync: vi.fn(),
+        });
+        const client = new DuckdbWarehouseClient_1.DuckdbWarehouseClient(undefined, {
+            instanceCacheKey: 'shared-query-hardening',
+            sharedResourceLimits: { memoryLimit: '256MB', threads: 1 },
+        });
+        await client.runQuery('SELECT 1 AS val');
+        expect(queryRunMock.mock.calls.map(([sql]) => sql)).toEqual([
+            "SET memory_limit = '256MB';",
+            'SET threads = 1;',
+            'SET allow_community_extensions = false;',
+            'SET autoinstall_known_extensions = false;',
+            'SET autoload_known_extensions = false;',
+            'SET allow_unredacted_secrets = false;',
+        ]);
+    });
+    it('confines embedded user queries without writing profiles to local files', async () => {
+        const runMock = vi.fn();
+        const streamMock = vi.fn(async () => getMockStreamResult([[{ val: 1 }]], [DUCKDB_TYPE_IDS.INTEGER]));
+        const logger = { info: vi.fn() };
+        createInstanceMock.mockResolvedValue(createMockConnection(streamMock, runMock));
+        const client = new DuckdbWarehouseClient_1.DuckdbWarehouseClient({
+            type: common_1.WarehouseTypes.DUCKDB,
+            connectionType: common_1.DuckdbConnectionType.EMBEDDED,
+            dataset: 'jaffle_shop',
+        }, { logger });
+        await client.runQuery('SELECT 1 AS val');
+        expect(createInstanceMock).toHaveBeenCalledWith(expect.stringContaining('jaffle_shop.duckdb'), {
+            access_mode: 'READ_ONLY',
+            memory_limit: '256MB',
+            threads: '1',
+        });
+        expect(runMock).toHaveBeenCalledWith("SET disabled_filesystems = 'LocalFileSystem';");
+        expect(runMock).not.toHaveBeenCalledWith("PRAGMA enable_profiling='json';");
+    });
+    it('rejects embedded queries beyond the process concurrency budget', async () => {
+        const pendingStreams = [];
+        const streamMock = vi.fn(() => new Promise((resolve) => {
+            pendingStreams.push(resolve);
+        }));
+        createInstanceMock.mockImplementation(async () => createMockConnection(streamMock));
+        const createClient = () => new DuckdbWarehouseClient_1.DuckdbWarehouseClient({
+            type: common_1.WarehouseTypes.DUCKDB,
+            connectionType: common_1.DuckdbConnectionType.EMBEDDED,
+            dataset: 'jaffle_shop',
+        });
+        const firstQuery = createClient().runQuery('SELECT 1 AS val');
+        const secondQuery = createClient().runQuery('SELECT 2 AS val');
+        await vi.waitFor(() => expect(streamMock).toHaveBeenCalledTimes(2));
+        await expect(createClient().runQuery('SELECT 3 AS val')).rejects.toThrow('Playground query capacity is full. Try again shortly.');
+        expect(createInstanceMock).toHaveBeenCalledTimes(2);
+        pendingStreams.forEach((resolve) => resolve(getMockStreamResult([[{ val: 1 }]], [DUCKDB_TYPE_IDS.INTEGER])));
+        await Promise.all([firstQuery, secondQuery]);
+    });
+    it('rejects concurrent embedded queries for the same organization', async () => {
+        let resolveStream = () => { };
+        const streamMock = vi.fn(() => new Promise((resolve) => {
+            resolveStream = resolve;
+        }));
+        createInstanceMock.mockImplementation(async () => createMockConnection(streamMock));
+        const createClient = () => new DuckdbWarehouseClient_1.DuckdbWarehouseClient({
+            type: common_1.WarehouseTypes.DUCKDB,
+            connectionType: common_1.DuckdbConnectionType.EMBEDDED,
+            dataset: 'jaffle_shop',
+        });
+        const tags = { organization_uuid: 'organization-1' };
+        const firstQuery = createClient().runQuery('SELECT 1 AS val', tags);
+        await vi.waitFor(() => expect(streamMock).toHaveBeenCalledOnce());
+        await expect(createClient().runQuery('SELECT 2 AS val', tags)).rejects.toThrow('Playground query capacity is full. Try again shortly.');
+        resolveStream(getMockStreamResult([[{ val: 1 }]], [DUCKDB_TYPE_IDS.INTEGER]));
+        await firstQuery;
+    });
+    it('logs structured profiles for non-embedded queries and reports metrics', async () => {
         const runMock = vi.fn(async (sql) => {
             const match = sql.match(/^PRAGMA profiling_output='(.+)';$/);
             if (match) {
@@ -362,8 +455,12 @@ describe('DuckdbWarehouseClient', () => {
         });
         const streamMock = vi.fn(async () => getMockStreamResult([[{ val: 1 }]], [DUCKDB_TYPE_IDS.INTEGER]));
         const logger = { info: vi.fn() };
+        const onQueryProfile = vi.fn();
         createInstanceMock.mockResolvedValue(createMockConnection(streamMock, runMock));
-        const client = new DuckdbWarehouseClient_1.DuckdbWarehouseClient(undefined, { logger });
+        const client = new DuckdbWarehouseClient_1.DuckdbWarehouseClient(undefined, {
+            logger,
+            onQueryProfile,
+        });
         await client.runQuery('SELECT 1 AS val', {
             query_uuid: 'query-123',
             chart_uuid: 'chart-123',
@@ -384,8 +481,14 @@ describe('DuckdbWarehouseClient', () => {
             bytesRead: 20225287,
             scanAmplification: 9905024 / 68,
         }));
+        expect(onQueryProfile).toHaveBeenCalledWith(expect.objectContaining({
+            latencyMs: 4747,
+            readParquetMs: 4632,
+            bytesRead: 20225287,
+        }));
+        expect(runMock).not.toHaveBeenCalledWith("SET disabled_filesystems = 'LocalFileSystem';");
     });
-    it('should log raw profile timings when DuckDB reports cpu above latency', async () => {
+    it('logs raw profile timings when DuckDB reports cpu above latency', async () => {
         const runMock = vi.fn(async (sql) => {
             const match = sql.match(/^PRAGMA profiling_output='(.+)';$/);
             if (match) {
@@ -479,6 +582,11 @@ describe('DuckdbWarehouseClient', () => {
             'read_text',
             'read_blob',
             'read_xlsx',
+            'parquet_scan',
+            'glob',
+            'sqlite_scan',
+            'postgres_scan',
+            'mysql_scan',
         ])('should reject user queries with %s()', async (blockedFunction) => {
             const streamMock = vi.fn(async () => getMockStreamResult([[{ val: 1 }]], [DUCKDB_TYPE_IDS.INTEGER]));
             const extractStatementsMock = createMockExtractStatements();
@@ -706,10 +814,11 @@ describe('DuckdbWarehouseClient', () => {
             // No explicit INSTALL/LOAD in DuckLake mode — rely on autoload.
             expect(joined).not.toMatch(/INSTALL httpfs/);
             expect(joined).not.toMatch(/LOAD httpfs/);
-            // Hardening flips autoload to TRUE for DuckLake.
             expect(stmts).toEqual(expect.arrayContaining([
                 'SET autoinstall_known_extensions = true;',
                 'SET autoload_known_extensions = true;',
+                'SET autoinstall_known_extensions = false;',
+                'SET autoload_known_extensions = false;',
                 'SET allow_community_extensions = false;',
                 'SET allow_unredacted_secrets = false;',
             ]));
@@ -721,6 +830,8 @@ describe('DuckdbWarehouseClient', () => {
             expect(dataIdx).toBeGreaterThan(catalogIdx);
             expect(duckLakeSecretIdx).toBeGreaterThan(dataIdx);
             expect(attachIdx).toBeGreaterThan(duckLakeSecretIdx);
+            expect(stmts.lastIndexOf('SET autoinstall_known_extensions = false;')).toBeGreaterThan(attachIdx);
+            expect(stmts.lastIndexOf('SET autoload_known_extensions = false;')).toBeGreaterThan(attachIdx);
             expect(stmts[catalogIdx]).toMatch(/TYPE postgres/);
             expect(stmts[catalogIdx]).toMatch(/HOST 'pg.example.com'/);
             expect(stmts[catalogIdx]).toMatch(/PASSWORD 'p@ss'/);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lightdash/warehouses",
-  "version": "0.3453.0",
+  "version": "0.3454.0",
   "description": "Warehouse connectors for Lightdash",
   "license": "MIT",
   "repository": {
@@ -30,7 +30,7 @@
     "snowflake-sdk": "2.3.4",
     "ssh2": "1.14.0",
     "trino-client": "0.2.9",
-    "@lightdash/common": "0.3453.0"
+    "@lightdash/common": "0.3454.0"
   },
   "devDependencies": {
     "@types/big.js": "6.2.2",