diff --git a/agents/drivers/cassandra-go/integration_test.go b/agents/drivers/cassandra-go/integration_test.go index 0172808822..9f9efa6202 100644 --- a/agents/drivers/cassandra-go/integration_test.go +++ b/agents/drivers/cassandra-go/integration_test.go @@ -142,6 +142,144 @@ func TestCassandraIntegration(t *testing.T) { } } +func TestCassandraTableDataCursorIntegration(t *testing.T) { + host := strings.TrimSpace(os.Getenv("CASSANDRA_TEST_HOST")) + if host == "" { + t.Skip("Cassandra integration environment is not configured") + } + port, err := strconv.Atoi(envDefault("CASSANDRA_TEST_PORT", "9042")) + if err != nil { + t.Fatal(err) + } + ssl, err := strconv.ParseBool(envDefault("CASSANDRA_TEST_SSL", "false")) + if err != nil { + t.Fatal(err) + } + connection := connectParams{ + Host: host, + Port: port, + Username: os.Getenv("CASSANDRA_TEST_USERNAME"), + Password: os.Getenv("CASSANDRA_TEST_PASSWORD"), + URLParams: os.Getenv("CASSANDRA_TEST_URL_PARAMS"), + SSL: ssl, + CACertPath: os.Getenv("CASSANDRA_TEST_CA_CERT_PATH"), + ClientCertPath: os.Getenv("CASSANDRA_TEST_CLIENT_CERT_PATH"), + ClientKeyPath: os.Getenv("CASSANDRA_TEST_CLIENT_KEY_PATH"), + } + runtime, err := newConnectionRuntime(connection) + if err != nil { + t.Fatal(err) + } + t.Cleanup(runtime.close) + server := newServer(runtime, connection) + if err := server.validateConnection(); err != nil { + t.Fatal(err) + } + + keyspace := "dbx_issue7465_" + strconv.FormatInt(time.Now().UnixNano(), 36) + mustCQL(t, server, "CREATE KEYSPACE "+quoteCQLIdentifier(keyspace)+" WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}", "") + t.Cleanup(func() { + if _, err := server.executeQuery(queryOptions{SQL: "DROP KEYSPACE IF EXISTS " + quoteCQLIdentifier(keyspace)}); err != nil { + t.Errorf("drop integration keyspace %s: %v", keyspace, err) + } + }) + mustCQL(t, server, "CREATE TABLE "+qualifiedCQLName(keyspace, "paged_rows")+" (id int PRIMARY KEY)", keyspace) + mustCQL(t, server, "CREATE TABLE "+qualifiedCQLName(keyspace, "empty_rows")+" (id int PRIMARY KEY)", keyspace) + for start := 0; start < 205; start += 25 { + end := min(start+25, 205) + statements := make([]string, 0, end-start) + for id := start; id < end; id++ { + statements = append(statements, fmt.Sprintf("INSERT INTO %s (id) VALUES (%d)", qualifiedCQLName(keyspace, "paged_rows"), id)) + } + mustStatements(t, server, keyspace, statements, false) + } + + tableSQL := `SELECT * FROM "paged_rows";` + page, err := server.executeQueryPage(queryOptions{SQL: tableSQL, Schema: keyspace, MaxRows: 1_000}, 100) + if err != nil || len(page.Rows) != 100 || !page.HasMore || page.SessionID == nil { + t.Fatalf("first page mismatch: rows=%d has_more=%v session=%v err=%v", len(page.Rows), page.HasMore, page.SessionID, err) + } + firstSessionID := *page.SessionID + page, err = server.fetchQueryPage(firstSessionID, 100) + if err != nil || len(page.Rows) != 100 || !page.HasMore || page.SessionID == nil || *page.SessionID != firstSessionID { + t.Fatalf("middle page mismatch: rows=%d has_more=%v session=%v err=%v", len(page.Rows), page.HasMore, page.SessionID, err) + } + page, err = server.fetchQueryPage(firstSessionID, 100) + if err != nil || len(page.Rows) != 5 || page.HasMore || page.SessionID != nil { + t.Fatalf("last page mismatch: rows=%d has_more=%v session=%v err=%v", len(page.Rows), page.HasMore, page.SessionID, err) + } + if _, err := server.fetchQueryPage(firstSessionID, 100); err == nil { + t.Fatal("exhausted query session remained usable") + } + + page, err = server.executeQueryPage(queryOptions{SQL: tableSQL, Schema: keyspace, MaxRows: 150}, 100) + if err != nil || len(page.Rows) != 100 || page.SessionID == nil { + t.Fatalf("maxRows first page mismatch: page=%v err=%v", page, err) + } + page, err = server.fetchQueryPage(*page.SessionID, 100) + if err != nil || len(page.Rows) != 50 || page.HasMore || !page.Truncated { + t.Fatalf("maxRows terminal page mismatch: rows=%d has_more=%v truncated=%v err=%v", len(page.Rows), page.HasMore, page.Truncated, err) + } + + page, err = server.executeQueryPage(queryOptions{SQL: tableSQL, Schema: keyspace, MaxRows: 1_000}, 100) + if err != nil || page.SessionID == nil { + t.Fatalf("close-session setup failed: page=%v err=%v", page, err) + } + closedSessionID := *page.SessionID + if !server.closeQuerySession(closedSessionID) { + t.Fatal("query session was not closed") + } + if _, err := server.fetchQueryPage(closedSessionID, 100); err == nil { + t.Fatal("closed query session remained usable") + } + + page, err = server.executeQueryPage(queryOptions{SQL: tableSQL, Schema: keyspace, MaxRows: 1_000}, 40) + if err != nil || len(page.Rows) != 40 || page.SessionID == nil { + t.Fatalf("page-size restart mismatch: page=%v err=%v", page, err) + } + page, err = server.fetchQueryPage(*page.SessionID, 40) + if err != nil || len(page.Rows) != 40 { + t.Fatalf("page-size continuation mismatch: rows=%d err=%v", len(page.Rows), err) + } + if page.SessionID != nil { + server.closeQuerySession(*page.SessionID) + } + + empty, err := server.executeQueryPage(queryOptions{SQL: `SELECT * FROM "empty_rows";`, Schema: keyspace, MaxRows: 1_000}, 100) + if err != nil || len(empty.Rows) != 0 || empty.HasMore || empty.SessionID != nil { + t.Fatalf("empty page mismatch: page=%v err=%v", empty, err) + } + + deletes := make([]string, 0, 5) + for id := 200; id < 205; id++ { + deletes = append(deletes, fmt.Sprintf("DELETE FROM %s WHERE id = %d", qualifiedCQLName(keyspace, "paged_rows"), id)) + } + mustStatements(t, server, keyspace, deletes, false) + page, err = server.executeQueryPage(queryOptions{SQL: tableSQL, Schema: keyspace, MaxRows: 1_000}, 100) + if err != nil { + t.Fatal(err) + } + totalRows := 0 + pageCount := 0 + for { + totalRows += len(page.Rows) + pageCount++ + if !page.HasMore { + break + } + if page.SessionID == nil || pageCount >= 3 { + t.Fatalf("exact-boundary cursor did not terminate: page=%v page_count=%d", page, pageCount) + } + page, err = server.fetchQueryPage(*page.SessionID, 100) + if err != nil { + t.Fatal(err) + } + } + if totalRows != 200 { + t.Fatalf("exact-boundary row count mismatch: got %d", totalRows) + } +} + func TestCassandraMaterializedViewDDLIntegration(t *testing.T) { host := strings.TrimSpace(os.Getenv("CASSANDRA_TEST_HOST")) if host == "" { diff --git a/apps/desktop/src/components/objects/ObjectBrowser.vue b/apps/desktop/src/components/objects/ObjectBrowser.vue index eca325cf4f..28e8919475 100644 --- a/apps/desktop/src/components/objects/ObjectBrowser.vue +++ b/apps/desktop/src/components/objects/ObjectBrowser.vue @@ -2290,16 +2290,31 @@ function tableDdlObjectType(type: ObjectBrowserRow["type"]): ObjectSourceKind | } async function exportDataLegacy(row: ObjectBrowserRow, format: "json") { + const { connection, database, catalog } = props; + const { id: connectionId, db_type: databaseType, query_timeout_secs: timeoutSecs } = connection; + const exportDatabaseType = effectiveDatabaseType.value; try { const schema = row.schema || selectedSchema.value; - const queryColumns = props.connection.db_type === "neo4j" ? (await api.getColumns(props.connection.id, props.database, schema || props.database, row.name, props.catalog)).map((column) => column.name) : undefined; + const queryColumns = databaseType === "neo4j" ? (await api.getColumns(connectionId, database, schema || database, row.name, catalog)).map((column) => column.name) : undefined; + const useAgentCursor = databaseType === "cassandra"; + const clientSessionId = useAgentCursor ? `table-export:${generateDatabaseExportId()}` : undefined; const result = await fetchTableDataForExport({ - databaseType: effectiveDatabaseType.value, - identifierQuote: connectionStore.connectionIdentifierQuote?.(props.connection.id), + databaseType: exportDatabaseType, + identifierQuote: connectionStore.connectionIdentifierQuote?.(connectionId), schema, tableName: row.name, columns: queryColumns, - executePage: (sql) => api.executeQuery(props.connection.id, props.database, sql), + useAgentCursor, + executePage: (sql, cursorOptions) => (cursorOptions ? api.executeQuery(connectionId, database, sql, undefined, undefined, { ...cursorOptions, clientSessionId, catalog, timeoutSecs }) : api.executeQuery(connectionId, database, sql)), + closeCursor: useAgentCursor + ? async (sessionId) => { + try { + if (sessionId) await api.closeQuerySession(connectionId, database, sessionId, clientSessionId, catalog); + } finally { + await api.closeClientConnectionSession(connectionId, database, clientSessionId!, catalog); + } + } + : undefined, }); if (format === "json") { diff --git a/apps/desktop/src/composables/__tests__/useDataGridActions.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridActions.spec.ts index 0c74eb17fc..9eb6a5e911 100644 --- a/apps/desktop/src/composables/__tests__/useDataGridActions.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDataGridActions.spec.ts @@ -589,6 +589,70 @@ describe("useDataGridActions", () => { expect(mocks.executeTabSql).toHaveBeenCalledWith("tab-1", "SELECT * FROM public.users LIMIT 100 OFFSET 100", expect.objectContaining({ appendResult: { maxRows: 10_000 } })); }); + it("continues the matching Cassandra table cursor and client session", async () => { + mocks.infiniteScroll = false; + mocks.getConfig.mockReturnValue({ id: "cassandra-1", db_type: "cassandra" }); + mocks.buildTableSelectSql.mockResolvedValueOnce('SELECT * FROM "paged_rows";'); + const active = tableDataTab({ id: "tab-a" }); + const target = tableDataTab({ + id: "tab-b", + connectionId: "cassandra-1", + resultPageLimit: 100, + resultPageOffset: 0, + resultClientSessionId: "tab-b", + result: { + columns: ["id"], + rows: Array.from({ length: 100 }, (_, index) => [index]), + affected_rows: 0, + execution_time_ms: 1, + session_id: "cassandra-page-1", + has_more: true, + }, + }); + mocks.tabs.push(active, target); + const actions = useDataGridActions(computed(() => active)); + + await actions.onPaginate(target.id, 100, 100); + + expect(mocks.executeTabSql).toHaveBeenCalledWith( + "tab-b", + 'SELECT * FROM "paged_rows";', + expect.objectContaining({ + pagination: { offset: 100, limit: 100, sessionId: "cassandra-page-1", clientSessionId: "tab-b" }, + }), + ); + }); + + it("restarts Cassandra table cursors for jumps and page-size changes", async () => { + mocks.infiniteScroll = false; + mocks.getConfig.mockReturnValue({ id: "cassandra-1", db_type: "cassandra" }); + mocks.buildTableSelectSql.mockResolvedValue('SELECT * FROM "paged_rows";'); + const tab = tableDataTab({ + connectionId: "cassandra-1", + resultPageLimit: 100, + resultPageOffset: 0, + resultClientSessionId: "tab-1", + result: { + columns: ["id"], + rows: Array.from({ length: 100 }, (_, index) => [index]), + affected_rows: 0, + execution_time_ms: 1, + session_id: "cassandra-page-1", + has_more: true, + }, + }); + mocks.tabs.push(tab); + const actions = useDataGridActions(computed(() => tab)); + + await actions.onPaginate(tab.id, 300, 100); + await actions.onPaginate(tab.id, 100, 50); + + expect(mocks.executeTabSql.mock.calls.map((call) => call[2].pagination)).toEqual([ + { offset: 300, limit: 100, sessionId: undefined, clientSessionId: undefined }, + { offset: 100, limit: 50, sessionId: undefined, clientSessionId: undefined }, + ]); + }); + it("ignores a stale structured order when its column was renamed", async () => { const tab = tableDataTab({ resultSortColumn: "old_name", diff --git a/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts b/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts index e76b433bad..d3aad0cfee 100644 --- a/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts +++ b/apps/desktop/src/composables/__tests__/useSidebarTreeExportRuntime.spec.ts @@ -8,6 +8,8 @@ const addExportTaskMock = vi.hoisted(() => vi.fn()); const updateTableExportTaskMock = vi.hoisted(() => vi.fn()); const apiMock = vi.hoisted(() => ({ buildTableSelectSql: vi.fn(async () => 'SELECT * FROM "main"."users" LIMIT 10000'), + closeClientConnectionSession: vi.fn(), + closeQuerySession: vi.fn(), executeQuery: vi.fn(), exportQueryResultCsv: vi.fn(), exportQueryResultJson: vi.fn(), @@ -242,6 +244,70 @@ describe("useSidebarTreeExportRuntime", () => { expect(toastMock).toHaveBeenCalledWith("导出失败：上一个 DuckDB 查询仍在停止中，请稍后重试。", 5000); }); + it("isolates and advances Cassandra JSON export cursors", async () => { + apiMock.buildTableSelectSql.mockResolvedValue('SELECT * FROM "events";'); + apiMock.executeQuery.mockResolvedValueOnce({ columns: ["id"], rows: [[1]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-1", has_more: true }).mockResolvedValueOnce({ columns: ["id"], rows: [[2]], affected_rows: 0, execution_time_ms: 1, has_more: false }); + const activeNode = shallowRef({ id: "table-1", type: "table", label: "events", connectionId: "cassandra-1", database: "app", children: [] } as TreeNode); + const connectionStore = { + ensureConnected: vi.fn(), + getConfig: vi.fn(() => ({ db_type: "cassandra", query_timeout_secs: 30 })), + connectionIdentifierQuote: vi.fn(() => '"'), + treeNodes: [], + selectedTreeNodeIds: [], + }; + const runtime = useSidebarTreeExportRuntime({ + activeNode, + connectionStore: connectionStore as never, + settingsStore: exportSettings() as never, + acceptedSelectionIds: () => null, + }); + + await runtime.exportData("json"); + + const firstOptions = apiMock.executeQuery.mock.calls[0]?.[5]; + const secondOptions = apiMock.executeQuery.mock.calls[1]?.[5]; + expect(firstOptions).toMatchObject({ maxRows: 2_147_483_647, fetchSize: 10_000, pageSize: 10_000, resultSessionId: undefined, timeoutSecs: 30 }); + expect(secondOptions).toMatchObject({ maxRows: 2_147_483_647, fetchSize: 10_000, pageSize: 10_000, resultSessionId: "cursor-1", clientSessionId: firstOptions.clientSessionId, timeoutSecs: 30 }); + expect(firstOptions.clientSessionId).toMatch(/^table-export:/); + expect(apiMock.exportQueryResultJson).toHaveBeenCalledWith("events.json", ["id"], [[1], [2]]); + expect(apiMock.closeClientConnectionSession).toHaveBeenCalledWith("cassandra-1", "app", firstOptions.clientSessionId, undefined); + }); + + it("closes a failed Cassandra cursor before disposing its client session", async () => { + apiMock.buildTableSelectSql.mockResolvedValue('SELECT * FROM "events";'); + apiMock.executeQuery.mockResolvedValueOnce({ columns: ["id"], rows: [[1]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-1", has_more: true }).mockRejectedValueOnce(new Error("cursor lost")); + const closing = deferred(); + apiMock.closeQuerySession.mockReturnValueOnce(closing.promise); + const runtime = useSidebarTreeExportRuntime({ + activeNode: shallowRef({ id: "table-1", type: "table", label: "events", connectionId: "cassandra-1", database: "app", children: [] } as TreeNode), + connectionStore: { ensureConnected: vi.fn(), getConfig: () => ({ db_type: "cassandra" }), connectionIdentifierQuote: () => '"', treeNodes: [], selectedTreeNodeIds: [] } as never, + settingsStore: exportSettings() as never, + acceptedSelectionIds: () => null, + }); + const exported = runtime.exportData("json"); + await vi.waitFor(() => expect(apiMock.closeQuerySession).toHaveBeenCalled()); + expect(apiMock.closeClientConnectionSession).not.toHaveBeenCalled(); + closing.reject(new Error("close failed")); + await exported; + expect(apiMock.closeClientConnectionSession).toHaveBeenCalledTimes(1); + expect(apiMock.exportQueryResultJson).not.toHaveBeenCalled(); + expect(toastMock).toHaveBeenCalledWith("导出失败：cursor lost", 5000); + }); + + it("preserves SQL Server legacy JSON export without enabling a new cursor path", async () => { + apiMock.executeQuery.mockResolvedValueOnce({ columns: ["id"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }); + const runtime = useSidebarTreeExportRuntime({ + activeNode: shallowRef({ id: "table-1", type: "table", label: "events", connectionId: "sqlserver-1", database: "app", children: [] } as TreeNode), + connectionStore: { ensureConnected: vi.fn(), getConfig: () => ({ db_type: "sqlserver", driver_profile: "sqlserver-legacy" }), connectionIdentifierQuote: () => '"', treeNodes: [], selectedTreeNodeIds: [] } as never, + settingsStore: exportSettings() as never, + acceptedSelectionIds: () => null, + }); + await runtime.exportData("json"); + expect(apiMock.executeQuery.mock.calls[0]).toHaveLength(3); + expect(apiMock.closeClientConnectionSession).not.toHaveBeenCalled(); + expect(apiMock.exportQueryResultJson).toHaveBeenCalledWith("events.json", ["id"], [[1]]); + }); + it("exports a mongo collection through the save-file path without a setup dialog", async () => { apiMock.exportMongodbQuery.mockImplementation(async (_request, onProgress) => { onProgress({ exportId: "export-1", status: "done", documentsRead: 3, bytesWritten: 12, elapsedMs: 4 }); diff --git a/apps/desktop/src/composables/useDataGridActions.ts b/apps/desktop/src/composables/useDataGridActions.ts index 0dcdf5ee44..3d0373fc0b 100644 --- a/apps/desktop/src/composables/useDataGridActions.ts +++ b/apps/desktop/src/composables/useDataGridActions.ts @@ -23,6 +23,7 @@ import type { DataGridReloadIntent } from "@/lib/dataGrid/dataGridToolbar"; import { continuousQueryResultMaxRows } from "@/lib/dataGrid/queryResultRowLimit"; import { queryResultBaseSql, queryResultExecutionSql } from "@/lib/tabs/tabPresentation"; import { sqlExecutionTargetCapabilities } from "@/lib/database/sqlExecutionTargetCapabilities"; +import { usesAgentCursorForTableData } from "@/lib/database/databaseDriverManifest"; const DATA_TAB_METADATA_TTL_MS = TABLE_METADATA_CACHE_TTL_MS; @@ -602,8 +603,7 @@ export function useDataGridActions(activeTab: ComputedRef) const expectedNextOffset = appendResult ? tab.result?.rows.length : (tab.resultPageOffset ?? 0) + (tab.resultPageLimit ?? limit); const continuesResultSession = offset === expectedNextOffset && limit === tab.resultPageLimit; const connection = useConnectionStore().getConfig(tab.connectionId); - const isSqlServerLegacy = connection?.db_type === "sqlserver" && connection.driver_profile?.trim().toLowerCase() === "sqlserver-legacy"; - const sessionId = isSqlServerLegacy && tab.result?.has_more && tab.result.session_id && continuesResultSession ? tab.result.session_id : undefined; + const sessionId = usesAgentCursorForTableData(connection?.db_type, connection?.driver_profile) && tab.result?.has_more && tab.result.session_id && continuesResultSession ? tab.result.session_id : undefined; await queryStore.executeTabSql(tab.id, sql, { pagination: { offset, limit, sessionId, clientSessionId: sessionId ? tab.resultClientSessionId : undefined }, ...appendOptions, diff --git a/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts b/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts index a53eb09ac0..231506423c 100644 --- a/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts +++ b/apps/desktop/src/composables/useSidebarTreeExportRuntime.ts @@ -33,6 +33,7 @@ import { } from "@/components/sidebar/sidebarTreeDialogState"; import type { CsvQuoteMode } from "@/lib/export/csvQuoteMode"; import type { SqlInsertDialect } from "@/lib/export/sqlInsertMode"; +import { uuid } from "@/lib/common/utils"; type StructureCopyFormat = "tsv" | "markdown"; @@ -355,6 +356,8 @@ export function useSidebarTreeExportRuntime(options: SidebarTreeExportRuntimeOpt try { await connectionStore.ensureConnected(connectionId); const queryColumns = config.db_type === "neo4j" ? (await api.getColumns(connectionId, database, target.metadataSchema, target.tableName, target.catalog)).map((column) => column.name) : undefined; + const useAgentCursor = config.db_type === "cassandra"; + const clientSessionId = useAgentCursor ? `table-export:${uuid()}` : undefined; const result = await fetchTableDataForExport({ databaseType: target.databaseType, identifierQuote: target.identifierQuote, @@ -362,7 +365,17 @@ export function useSidebarTreeExportRuntime(options: SidebarTreeExportRuntimeOpt tableName: target.tableName, tableType: target.tableType, columns: queryColumns, - executePage: (sql) => api.executeQuery(connectionId, database, sql), + useAgentCursor, + executePage: (sql, cursorOptions) => (cursorOptions ? api.executeQuery(connectionId, database, sql, undefined, undefined, { ...cursorOptions, clientSessionId, catalog: target.catalog, timeoutSecs: config.query_timeout_secs }) : api.executeQuery(connectionId, database, sql)), + closeCursor: useAgentCursor + ? async (sessionId) => { + try { + if (sessionId) await api.closeQuerySession(connectionId, database, sessionId, clientSessionId, target.catalog); + } finally { + await api.closeClientConnectionSession(connectionId, database, clientSessionId!, target.catalog); + } + } + : undefined, }); const outputPath = await resolveTableExportOutputPath(target, "json", outputDirectory); diff --git a/apps/desktop/src/lib/__tests__/database/databaseDriverManifest.spec.ts b/apps/desktop/src/lib/__tests__/database/databaseDriverManifest.spec.ts index 6a38fe146a..3b59c33f87 100644 --- a/apps/desktop/src/lib/__tests__/database/databaseDriverManifest.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/databaseDriverManifest.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { databaseConnectionFormKind, databaseDefaultPort, databaseManifestEntry, databaseRuntimeMode, usesAgentCursorForQuery } from "@/lib/database/databaseDriverManifest"; +import { databaseConnectionFormKind, databaseDefaultPort, databaseManifestEntry, databaseRuntimeMode, usesAgentCursorForQuery, usesAgentCursorForTableData } from "@/lib/database/databaseDriverManifest"; describe("databaseDriverManifest", () => { it("uses agent cursor only for agent or external runtimes", () => { @@ -10,6 +10,7 @@ describe("databaseDriverManifest", () => { expect(usesAgentCursorForQuery("meilisearch")).toBe(false); expect(usesAgentCursorForQuery("dynamodb")).toBe(false); expect(usesAgentCursorForQuery("mysql")).toBe(false); + expect(usesAgentCursorForQuery("cassandra")).toBe(true); expect(usesAgentCursorForQuery("jdbc")).toBe(true); expect(usesAgentCursorForQuery("prestosql")).toBe(true); expect(usesAgentCursorForQuery("sqlserver")).toBe(false); @@ -17,6 +18,14 @@ describe("databaseDriverManifest", () => { expect(usesAgentCursorForQuery("sqlserver", " SQLSERVER-LEGACY ")).toBe(true); }); + it("limits table-data agent cursors to Cassandra and SQL Server legacy", () => { + expect(usesAgentCursorForTableData("cassandra")).toBe(true); + expect(usesAgentCursorForTableData("sqlserver", "sqlserver-legacy")).toBe(true); + expect(usesAgentCursorForTableData("sqlserver")).toBe(false); + expect(usesAgentCursorForTableData("jdbc")).toBe(false); + expect(usesAgentCursorForTableData("prestosql")).toBe(false); + }); + it("exposes connection defaults from the shared manifest", () => { expect(databaseDefaultPort("mysql")).toBe(3306); expect(databaseManifestEntry("h2")).toMatchObject({ diff --git a/apps/desktop/src/lib/database/databaseDriverManifest.ts b/apps/desktop/src/lib/database/databaseDriverManifest.ts index ebef349c8d..a9f6a12c1f 100644 --- a/apps/desktop/src/lib/database/databaseDriverManifest.ts +++ b/apps/desktop/src/lib/database/databaseDriverManifest.ts @@ -159,6 +159,10 @@ export function usesAgentCursorForQuery(dbType?: DatabaseType, driverProfile?: s return runtimeMode === "agent" || runtimeMode === "external"; } +export function usesAgentCursorForTableData(dbType?: DatabaseType, driverProfile?: string): boolean { + return dbType === "cassandra" || (dbType === "sqlserver" && driverProfile?.trim().toLowerCase() === "sqlserver-legacy"); +} + function productCapabilities(overrides: Partial): DatabaseProductCapabilities { const capabilities = Object.fromEntries(DATABASE_PRODUCT_CAPABILITY_KEYS.map((key) => [key, false])) as DatabaseProductCapabilities; return { ...capabilities, ...overrides }; diff --git a/apps/desktop/src/lib/table/tableDataExport.ts b/apps/desktop/src/lib/table/tableDataExport.ts index d631ffd6d3..a837dc2b97 100644 --- a/apps/desktop/src/lib/table/tableDataExport.ts +++ b/apps/desktop/src/lib/table/tableDataExport.ts @@ -2,6 +2,14 @@ import type { DatabaseType, QueryResult } from "@/types/database.ts"; import { buildTableSelectSql } from "@/lib/table/tableSelectSql.ts"; export const TABLE_DATA_EXPORT_PAGE_SIZE = 10_000; +export const TABLE_DATA_EXPORT_MAX_ROWS = 2_147_483_647; + +export interface TableDataExportCursorPageOptions { + maxRows: number; + fetchSize: number; + pageSize: number; + resultSessionId?: string; +} export interface FetchTableDataForExportOptions { databaseType?: DatabaseType; @@ -11,8 +19,10 @@ export interface FetchTableDataForExportOptions { tableType?: string; columns?: string[]; pageSize?: number; + useAgentCursor?: boolean; buildPageSql?: (options: { databaseType?: DatabaseType; schema?: string; tableName: string; tableType?: string; columns?: string[]; limit: number; offset: number }) => Promise | string; - executePage: (sql: string) => Promise; + executePage: (sql: string, cursorOptions?: TableDataExportCursorPageOptions) => Promise; + closeCursor?: (sessionId?: string) => Promise; } export async function fetchTableDataForExport(options: FetchTableDataForExportOptions): Promise { @@ -33,35 +43,60 @@ export async function fetchTableDataForExport(options: FetchTableDataForExportOp const rows: QueryResult["rows"] = []; let columns: string[] = []; let executionTimeMs = 0; + let sessionId: string | undefined; - while (true) { - const sql = await (options.buildPageSql ?? buildTableSelectSql)({ - databaseType: options.databaseType, - identifierQuote: options.identifierQuote, - schema: options.schema, - tableName: options.tableName, - tableType: options.tableType, - columns: options.columns, - limit: pageSize, - offset, - }); - const result = await options.executePage(sql); - if (columns.length === 0) columns = result.columns; - rows.push(...result.rows); - executionTimeMs += result.execution_time_ms ?? 0; + try { + while (true) { + const sql = await (options.buildPageSql ?? buildTableSelectSql)({ + databaseType: options.databaseType, + identifierQuote: options.identifierQuote, + schema: options.schema, + tableName: options.tableName, + tableType: options.tableType, + columns: options.columns, + limit: pageSize, + offset, + }); + const result = await options.executePage( + sql, + options.useAgentCursor + ? { + maxRows: TABLE_DATA_EXPORT_MAX_ROWS, + fetchSize: pageSize, + pageSize, + resultSessionId: sessionId, + } + : undefined, + ); + const nextSessionId = result.session_id?.trim() || undefined; + if (options.useAgentCursor) { + sessionId = nextSessionId ?? sessionId; + if (result.truncated || typeof result.has_more !== "boolean") throw new Error("Incomplete cursor result during table export"); + if (result.has_more && !nextSessionId) throw new Error("Result session ended before table export completed"); + } + if (columns.length === 0) columns = result.columns; + rows.push(...result.rows); + executionTimeMs += result.execution_time_ms ?? 0; + sessionId = nextSessionId; - if (result.rows.length < pageSize) { - return { - columns, - rows, - affected_rows: 0, - execution_time_ms: executionTimeMs, - truncated: false, - session_id: undefined, - has_more: false, - }; + if (options.useAgentCursor ? result.has_more !== true : result.rows.length < pageSize) { + return { + columns, + rows, + affected_rows: 0, + execution_time_ms: executionTimeMs, + truncated: false, + session_id: undefined, + has_more: false, + }; + } + offset += result.rows.length; + } + } finally { + if (options.useAgentCursor && options.closeCursor) { + try { + await options.closeCursor(sessionId); + } catch {} } - - offset += result.rows.length; } } diff --git a/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts index ef8b1b1e1f..b4eacb80b5 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts @@ -248,6 +248,227 @@ describe("queryStore table data refresh", () => { expect(mocks.executeMulti).toHaveBeenCalledWith("yasdb-1", "app", "SELECT * FROM app.users", undefined, expect.any(String), expect.not.objectContaining({ rowOffset: expect.anything() })); }); + it("continues and restarts Cassandra table cursors without duplicating page one", async () => { + mocks.getConnectionConfig.mockReturnValue({ + id: "cassandra-1", + name: "Cassandra", + db_type: "cassandra", + database: "app", + query_timeout_secs: 30, + }); + mocks.executeMulti + .mockResolvedValueOnce([{ columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-1", has_more: true }]) + .mockResolvedValueOnce([{ columns: ["id"], rows: [[3], [4]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-1", has_more: true }]) + .mockResolvedValueOnce([{ columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-2", has_more: true }]) + .mockResolvedValueOnce([{ columns: ["id"], rows: [[3], [4]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-2", has_more: true }]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("cassandra-1", "app", "paged_rows", "data"); + store.setTableMeta(tabId, { + tableName: "paged_rows", + tableType: "TABLE", + columns: [{ name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }], + primaryKeys: ["id"], + }); + const sql = 'SELECT * FROM "paged_rows";'; + + await store.executeTabSql(tabId, sql, { pagination: { limit: 2, offset: 0 } }); + let tab = store.tabs.find((candidate) => candidate.id === tabId)!; + expect(tab.result?.rows).toEqual([[1], [2]]); + expect(tab.resultSessionId).toBe("cursor-1"); + expect(tab.resultClientSessionId).toBe(tabId); + + await store.executeTabSql(tabId, sql, { + pagination: { limit: 2, offset: 2, sessionId: "cursor-1", clientSessionId: tab.resultClientSessionId }, + }); + tab = store.tabs.find((candidate) => candidate.id === tabId)!; + expect(tab.result?.rows).toEqual([[3], [4]]); + expect(mocks.closeQuerySession).not.toHaveBeenCalled(); + + await store.executeTabSql(tabId, sql, { pagination: { limit: 2, offset: 0 } }); + tab = store.tabs.find((candidate) => candidate.id === tabId)!; + expect(tab.result?.rows).toEqual([[1], [2]]); + expect(mocks.closeQuerySession).toHaveBeenCalledWith("cassandra-1", "app", "cursor-1", tabId); + + await store.executeTabSql(tabId, sql, { + pagination: { limit: 2, offset: 2, sessionId: "cursor-2", clientSessionId: tab.resultClientSessionId }, + }); + expect(store.tabs.find((candidate) => candidate.id === tabId)?.result?.rows).toEqual([[3], [4]]); + expect(mocks.executeMulti.mock.calls.map((call) => call[5])).toEqual([ + expect.objectContaining({ maxRows: 2_147_483_647, fetchSize: 2, pageSize: 2, resultSessionId: undefined, clientSessionId: tabId }), + expect.objectContaining({ maxRows: 2_147_483_647, fetchSize: 2, pageSize: 2, resultSessionId: "cursor-1", clientSessionId: tabId }), + expect.objectContaining({ maxRows: 2_147_483_647, fetchSize: 2, pageSize: 2, resultSessionId: undefined, clientSessionId: tabId }), + expect.objectContaining({ maxRows: 2_147_483_647, fetchSize: 2, pageSize: 2, resultSessionId: "cursor-2", clientSessionId: tabId }), + ]); + }); + + it("replays Cassandra table cursors to direct, last, and exact-boundary pages", async () => { + mocks.getConnectionConfig.mockReturnValue({ id: "cassandra-1", name: "Cassandra", db_type: "cassandra", database: "app" }); + const page = (rows: number[], sessionId?: string) => ({ + columns: ["id"], + rows: rows.map((value) => [value]), + affected_rows: 0, + execution_time_ms: rows.length, + session_id: sessionId, + has_more: !!sessionId, + query_timings_ms: { agent_total: rows.length }, + }); + mocks.executeMulti + .mockResolvedValueOnce([page([1, 2], "cursor-last")]) + .mockResolvedValueOnce([page([3, 4], "cursor-last")]) + .mockResolvedValueOnce([page([5])]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("cassandra-1", "app", "paged_rows", "data"); + store.setTableMeta(tabId, { + tableName: "paged_rows", + tableType: "TABLE", + columns: [{ name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }], + primaryKeys: ["id"], + }); + const sql = 'SELECT * FROM "paged_rows";'; + + await store.executeTabSql(tabId, sql, { pagination: { limit: 2, offset: 4 } }); + + let tab = store.tabs.find((candidate) => candidate.id === tabId)!; + expect(tab.result?.rows).toEqual([[5]]); + expect(tab.result?.execution_time_ms).toBe(5); + expect(tab.result?.timing_page_count).toBe(3); + expect(tab.resultPageOffset).toBe(4); + expect(tab.resultTotalRowCount).toBe(5); + expect(mocks.executeMulti.mock.calls.map((call) => call[5]?.resultSessionId)).toEqual([undefined, "cursor-last", "cursor-last"]); + + mocks.executeMulti.mockReset(); + mocks.executeMulti.mockResolvedValueOnce([page([1, 2], "cursor-boundary")]).mockResolvedValueOnce([page([3, 4])]); + await store.executeTabSql(tabId, sql, { pagination: { limit: 2, offset: 4 } }); + + tab = store.tabs.find((candidate) => candidate.id === tabId)!; + expect(tab.result?.rows).toEqual([]); + expect(tab.result?.has_more).toBe(false); + expect(tab.resultPageOffset).toBe(4); + expect(tab.resultTotalRowCount).toBeUndefined(); + expect(mocks.executeMulti).toHaveBeenCalledTimes(2); + }); + + it("closes a Cassandra replay cursor when a later page fails", async () => { + mocks.getConnectionConfig.mockReturnValue({ id: "cassandra-1", name: "Cassandra", db_type: "cassandra", database: "app" }); + mocks.executeMulti.mockResolvedValueOnce([{ columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-lost", has_more: true }]).mockRejectedValueOnce(new Error("query session not found: cursor-lost")); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("cassandra-1", "app", "paged_rows", "data"); + store.setTableMeta(tabId, { + tableName: "paged_rows", + tableType: "TABLE", + columns: [{ name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }], + primaryKeys: ["id"], + }); + + await store.executeTabSql(tabId, 'SELECT * FROM "paged_rows";', { pagination: { limit: 2, offset: 2 } }); + + const tab = store.tabs.find((candidate) => candidate.id === tabId)!; + expect(tab.result?.execution_error).toBe(true); + expect(tab.resultSessionId).toBeUndefined(); + expect(mocks.closeQuerySession).toHaveBeenCalledWith("cassandra-1", "app", "cursor-lost", tabId, undefined); + }); + + it.each(["cancellation", "supersession"] as const)("does not dispatch another Cassandra replay page after %s", async (interruption) => { + mocks.getConnectionConfig.mockReturnValue({ id: "cassandra-1", name: "Cassandra", db_type: "cassandra", database: "app" }); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("cassandra-1", "app", "paged_rows", "data"); + store.setTableMeta(tabId, { + tableName: "paged_rows", + tableType: "TABLE", + columns: [{ name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }], + primaryKeys: ["id"], + }); + mocks.executeMulti.mockImplementationOnce(async () => { + const tab = store.tabs.find((candidate) => candidate.id === tabId)!; + if (interruption === "cancellation") { + tab.isCancelling = true; + tab.cancelRequestCount = (tab.cancelRequestCount ?? 0) + 1; + } else { + tab.executionId = "replacement-execution"; + } + return [{ columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 1, session_id: "cursor-cancel", has_more: true }]; + }); + + await store.executeTabSql(tabId, 'SELECT * FROM "paged_rows";', { pagination: { limit: 2, offset: 4 } }); + + expect(mocks.executeMulti).toHaveBeenCalledTimes(1); + expect(mocks.closeQuerySession).toHaveBeenCalledWith("cassandra-1", "app", "cursor-cancel", tabId, undefined); + }); + + it("exports Cassandra table data by advancing one Agent cursor", async () => { + mocks.getConnectionConfig.mockReturnValue({ id: "cassandra-1", name: "Cassandra", db_type: "cassandra", database: "app" }); + mocks.buildTableSelectSql.mockResolvedValue('SELECT * FROM "paged_rows";'); + mocks.executeMulti.mockResolvedValueOnce([{ columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 2, session_id: "export-cursor", has_more: true }]).mockResolvedValueOnce([{ columns: ["id"], rows: [[3]], affected_rows: 0, execution_time_ms: 1, has_more: false }]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("cassandra-1", "app", "paged_rows", "data"); + store.setTableMeta(tabId, { + tableName: "paged_rows", + tableType: "TABLE", + columns: [{ name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }], + primaryKeys: ["id"], + }); + const tab = store.tabs.find((candidate) => candidate.id === tabId)!; + tab.result = { columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 2 }; + + const exported = await store.fetchTabResultForExport(tabId); + + expect(exported?.rows).toEqual([[1], [2], [3]]); + expect(exported?.execution_time_ms).toBe(3); + expect(mocks.buildTableSelectSql).toHaveBeenCalledTimes(2); + expect(mocks.executeMulti.mock.calls.map((call) => call[5])).toEqual([ + expect.objectContaining({ maxRows: 2_147_483_647, fetchSize: 10_000, pageSize: 10_000, resultSessionId: undefined, clientSessionId: `${tabId}:export` }), + expect.objectContaining({ maxRows: 2_147_483_647, fetchSize: 10_000, pageSize: 10_000, resultSessionId: "export-cursor", clientSessionId: `${tabId}:export` }), + ]); + expect(mocks.closeClientConnectionSession).toHaveBeenCalledWith("cassandra-1", "app", `${tabId}:export`); + }); + + it("stops Cassandra table export when continuation has no session", async () => { + mocks.getConnectionConfig.mockReturnValue({ id: "cassandra-1", name: "Cassandra", db_type: "cassandra", database: "app" }); + mocks.buildTableSelectSql.mockResolvedValue('SELECT * FROM "paged_rows";'); + mocks.executeMulti.mockResolvedValueOnce([{ columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 1, has_more: true }]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("cassandra-1", "app", "paged_rows", "data"); + store.setTableMeta(tabId, { + tableName: "paged_rows", + tableType: "TABLE", + columns: [{ name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }], + primaryKeys: ["id"], + }); + const tab = store.tabs.find((candidate) => candidate.id === tabId)!; + tab.result = { columns: ["id"], rows: [[1], [2]], affected_rows: 0, execution_time_ms: 1 }; + + await expect(store.fetchTabResultForExport(tabId)).rejects.toThrow("Result session ended before table export completed"); + expect(mocks.executeMulti).toHaveBeenCalledTimes(1); + expect(mocks.closeClientConnectionSession).toHaveBeenCalledWith("cassandra-1", "app", `${tabId}:export`); + }); + + it.each([{ has_more: undefined }, { has_more: false, truncated: true }, { has_more: true, session_id: " " }])("rejects incomplete Cassandra export continuation %j and cleans its active cursor", async (malformed) => { + mocks.getConnectionConfig.mockReturnValue({ id: "cassandra-1", name: "Cassandra", db_type: "cassandra", database: "app" }); + mocks.buildTableSelectSql.mockResolvedValue('SELECT * FROM "paged_rows";'); + mocks.executeMulti.mockResolvedValueOnce([{ columns: ["id"], rows: [[1]], affected_rows: 0, execution_time_ms: 1, session_id: "active-cursor", has_more: true }]).mockResolvedValueOnce([{ columns: ["id"], rows: [[2]], affected_rows: 0, execution_time_ms: 1, ...malformed }]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("cassandra-1", "app", "paged_rows", "data"); + store.setTableMeta(tabId, { + tableName: "paged_rows", + tableType: "TABLE", + columns: [{ name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }], + primaryKeys: ["id"], + }); + const tab = store.tabs.find((candidate) => candidate.id === tabId)!; + tab.result = { columns: ["id"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }; + await expect(store.fetchTabResultForExport(tabId)).rejects.toThrow(/export/); + expect(mocks.executeMulti).toHaveBeenCalledTimes(2); + expect(mocks.closeQuerySession).toHaveBeenCalledWith("cassandra-1", "app", "active-cursor", `${tabId}:export`, undefined); + expect(mocks.closeClientConnectionSession).toHaveBeenCalledWith("cassandra-1", "app", `${tabId}:export`); + }); + it("executes Doris external catalog data tabs against the catalog database, not the connection default database", async () => { mocks.getConnectionConfig.mockReturnValue({ id: "doris-1", diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 03d8c49127..087d94bc8a 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -62,7 +62,7 @@ import { redisCommandResultToQueryResult } from "@/lib/redis/redisQueryResult"; import { nextRedisCommandDb } from "@/lib/redis/redisCommandSession"; import { isRedisMutatingCommand } from "@/lib/redis/redisCommandTable"; import { formatRedisConsoleValue } from "@/lib/redis/redisValuePresentation"; -import { usesAgentCursorForQuery } from "@/lib/database/databaseDriverManifest"; +import { usesAgentCursorForQuery, usesAgentCursorForTableData } from "@/lib/database/databaseDriverManifest"; import { connectionIsDorisFamilyCatalogCapable, defaultAutoCommitForDbType, supportsClearableQuerySchema, supportsTransaction, usesOracleStickyTransactionState, usesProvenReadOnlyStickyTransactionState } from "@/lib/database/databaseFeatureSupport"; import { canInsertTableRows, canUseKeylessRowPredicate, DBX_ROWID_COLUMN, editablePrimaryKeys, shouldIncludeSyntheticRowId, usesSyntheticRowIdKey } from "@/lib/table/tableEditing"; import { TABLE_DATA_EXPORT_PAGE_SIZE } from "@/lib/table/tableDataExport"; @@ -6732,6 +6732,7 @@ export const useQueryStore = defineStore("query", () => { const traceId = executionId.slice(0, 8); const startedAt = performance.now(); const executionTargetEpoch = manualTransactionTargetEpoch(tab); + const cancelRequestCountAtStart = tab.cancelRequestCount ?? 0; const elapsed = () => `${Math.round(performance.now() - startedAt)}ms`; const batchResume = options?.batchResume; const continueOnBatchError = batchResume?.continueOnError ?? settingsStore.editorSettings.continueOnErrorOnBatch; @@ -7602,7 +7603,7 @@ export const useQueryStore = defineStore("query", () => { const pagination = limitQueryPagination(requestedPagination, queryResultMaxRows); pageLimit = pagination.limit; pageOffset = pagination.offset; - useAgentResultSession = conn?.db_type === "sqlserver" && conn?.driver_profile?.trim().toLowerCase() === "sqlserver-legacy"; + useAgentResultSession = usesAgentCursorForTableData(conn?.db_type, conn?.driver_profile); } errorLocateContext = { @@ -7670,7 +7671,7 @@ export const useQueryStore = defineStore("query", () => { }); executionDispatched = true; if (tab.mode === "query") clientRequestStartedAt = performance.now(); - if (useAgentResultSession && tab.mode === "query" && typeof pageOffset === "number" && pageOffset > 0 && !options?.pagination?.sessionId && !(tab.batchSqlExecution && tab.batchSqlExecution.total > 1)) { + if (useAgentResultSession && (tab.mode === "query" || tab.mode === "data") && typeof pageOffset === "number" && pageOffset > 0 && !options?.pagination?.sessionId && !(tab.batchSqlExecution && tab.batchSqlExecution.total > 1)) { return (async () => { let sessionId: string | undefined; let skipped = 0; @@ -7678,35 +7679,50 @@ export const useQueryStore = defineStore("query", () => { let executionMs = 0; let completeTimings = true; const timings: Record = {}; - while (true) { - const pageResults = await api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, { - ...executionOptions, - resultSessionId: sessionId, - }); - const page = pageResults[0]; - if (!page) return pageResults; - pageCount += 1; - executionMs += page.execution_time_ms; - if (!page.query_timings_ms) completeTimings = false; - else for (const [key, value] of Object.entries(page.query_timings_ms)) timings[key] = (timings[key] ?? 0) + value; - // Offset jumps consume several cursor pages before publication. - // Keep their timings as well, without retaining skipped rows. - const timingSummary = { query_timings_ms: completeTimings ? { ...timings } : undefined, execution_time_ms: executionMs, timing_page_count: pageCount }; - if (skipped + page.rows.length > pageOffset) { - const start = pageOffset - skipped; - const limit = typeof pageLimit === "number" ? pageLimit : page.rows.length - start; - return [{ ...page, ...timingSummary, rows: page.rows.slice(start, start + limit) }]; - } - skipped += page.rows.length; - if (!page.has_more || !page.session_id) { - if (page.has_more) { - // The cursor session ended before the requested offset; - // returning the short page would show the wrong rows. - throw new Error("Result session ended before the requested page offset"); + try { + while (true) { + const pageResults = await api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, { + ...executionOptions, + resultSessionId: sessionId, + }); + const page = pageResults[0]; + if (!page) return pageResults; + pageCount += 1; + executionMs += page.execution_time_ms; + if (!page.query_timings_ms) completeTimings = false; + else for (const [key, value] of Object.entries(page.query_timings_ms)) timings[key] = (timings[key] ?? 0) + value; + // Offset jumps consume several cursor pages before publication. + // Keep their timings as well, without retaining skipped rows. + const timingSummary = { query_timings_ms: completeTimings ? { ...timings } : undefined, execution_time_ms: executionMs, timing_page_count: pageCount }; + if (skipped + page.rows.length > pageOffset) { + const start = pageOffset - skipped; + const limit = typeof pageLimit === "number" ? pageLimit : page.rows.length - start; + return [{ ...page, ...timingSummary, rows: page.rows.slice(start, start + limit) }]; + } + skipped += page.rows.length; + if (!page.has_more || !page.session_id) { + if (page.has_more) { + // The cursor session ended before the requested offset; + // returning the short page would show the wrong rows. + throw new Error("Result session ended before the requested page offset"); + } + return [{ ...page, ...timingSummary, rows: [] }]; + } + sessionId = page.session_id; + const replayTab = findExecutionTab(id); + if (replayTab !== tab || replayTab.executionId !== executionId || replayTab.isCancelling || (replayTab.cancelRequestCount ?? 0) !== cancelRequestCountAtStart || manualTransactionTargetEpoch(replayTab) !== executionTargetEpoch) { + throw new Error("Query canceled"); } - return [{ ...page, ...timingSummary, rows: [] }]; } - sessionId = page.session_id; + } catch (error) { + if (sessionId) { + try { + await api.closeQuerySession(executionConnectionId, executionDatabase, sessionId, executionClientSessionId, executionCatalog); + } catch (closeError) { + queryExecutionLog("warn", "offset-replay-session-close:error", { traceId, sessionId, error: closeError }); + } + } + throw error; } })(); } @@ -8937,7 +8953,7 @@ export const useQueryStore = defineStore("query", () => { let offset = 0; const clientSessionId = tabClientSessionId(tab, "export"); const exportExecutionId = uuid(); - const useAgentCursor = conn?.db_type === "sqlserver" && conn?.driver_profile?.trim().toLowerCase() === "sqlserver-legacy"; + const useAgentCursor = usesAgentCursorForTableData(conn?.db_type, conn?.driver_profile); let sessionId: string | undefined; try { @@ -8983,19 +8999,33 @@ export const useQueryStore = defineStore("query", () => { }, ); const result = results[0]; - if (!result) break; + if (!result) { + if (useAgentCursor) throw new Error("Missing cursor result during table export"); + break; + } + const nextSessionId = result.session_id?.trim() || undefined; + if (useAgentCursor) { + sessionId = nextSessionId ?? sessionId; + if (conn?.db_type === "cassandra" && (result.truncated || typeof result.has_more !== "boolean")) throw new Error("Incomplete cursor result during table export"); + if (result.has_more === true && !nextSessionId) throw new Error("Result session ended before table export completed"); + } if (columns.length === 0) columns = result.columns; rows.push(...result.rows); executionTimeMs += result.execution_time_ms ?? 0; onProgress?.({ rowsExported: rows.length, totalRows }); - sessionId = result.session_id ?? undefined; + sessionId = nextSessionId; const shouldFetchNextPage = useAgentCursor ? result.has_more === true : result.rows.length >= pageLimit; if (!shouldFetchNextPage) break; offset += result.rows.length; } } finally { - if (sessionId) void api.closeQuerySession(tab.connectionId, executionDatabase, sessionId, clientSessionId, tableMeta.catalog); - void closeClientSessionId(tab.connectionId, executionDatabase, clientSessionId, tableMeta.catalog, { tabId: tab.id }); + try { + if (sessionId) await api.closeQuerySession(tab.connectionId, executionDatabase, sessionId, clientSessionId, tableMeta.catalog); + } catch (error) { + queryExecutionLog("warn", "table-export-session-close:error", { sessionId, error }); + } finally { + await closeClientSessionId(tab.connectionId, executionDatabase, clientSessionId, tableMeta.catalog, { tabId: tab.id }); + } } return { diff --git a/crates/dbx-core/src/data/csv_export.rs b/crates/dbx-core/src/data/csv_export.rs index 6c841dc49f..c410a429d9 100644 --- a/crates/dbx-core/src/data/csv_export.rs +++ b/crates/dbx-core/src/data/csv_export.rs @@ -1,11 +1,13 @@ use serde::{Deserialize, Serialize}; use std::fs::File; +use std::future::Future; use std::io::{BufWriter, Write}; use crate::connection::AppState; use crate::models::connection::DatabaseType; use crate::query::{execute_sql_statement_with_options, QueryExecutionOptions}; use crate::sql_dialect::{build_table_data_select_sql, TableDataSelectSqlOptions}; +use crate::types::QueryResult; pub use dbx_formats::csv_export::*; @@ -52,10 +54,53 @@ pub async fn export_table_data_csv_core(state: &AppState, options: TableCsvExpor if database_type == DatabaseType::Salesforce { return Err("Exporting a Salesforce object to CSV is not supported yet: SOQL pages through a QueryLocator, not LIMIT/OFFSET. Run a SOQL query and export its result instead.".to_string()); } - let page_size = options.page_size.unwrap_or(TABLE_DATA_EXPORT_PAGE_SIZE).max(1); let mut writer = BufWriter::new(File::create(&options.file_path).map_err(|err| format!("Failed to write CSV file: {err}"))?); writer.write_all("\u{FEFF}".as_bytes()).map_err(|err| err.to_string())?; + let client_session_id = + (database_type == DatabaseType::Cassandra).then(|| format!("table-export:{}", uuid::Uuid::new_v4())); + let export_options = &options; + let outcome = write_table_csv_pages( + &mut writer, + database_type, + &options, + client_session_id.as_deref(), + |sql, query_options| async move { + execute_sql_statement_with_options( + state, + &export_options.connection_id, + &export_options.database, + &sql, + export_options.schema.as_deref(), + None, + query_options, + ) + .await + }, + ) + .await; + if let Some(client_session_id) = client_session_id { + let _ = + state.close_client_session_pool(&options.connection_id, Some(&options.database), &client_session_id).await; + } + outcome +} + +async fn write_table_csv_pages( + writer: &mut impl Write, + database_type: DatabaseType, + options: &TableCsvExportOptions, + client_session_id: Option<&str>, + mut execute_page: Execute, +) -> Result +where + Execute: FnMut(String, QueryExecutionOptions) -> QueryFuture, + QueryFuture: Future>, +{ + let use_cursor = database_type == DatabaseType::Cassandra; + let requested_page_size = options.page_size.unwrap_or(TABLE_DATA_EXPORT_PAGE_SIZE).max(1); + let page_size = if use_cursor { requested_page_size.min(TABLE_DATA_EXPORT_PAGE_SIZE) } else { requested_page_size }; + let mut session_id = None; let mut offset = 0usize; let mut rows_exported = 0u64; @@ -77,37 +122,45 @@ pub async fn export_table_data_csv_core(state: &AppState, options: TableCsvExpor include_row_id: false, ..Default::default() }); - let result = execute_sql_statement_with_options( - state, - &options.connection_id, - &options.database, - &sql, - options.schema.as_deref(), - None, + let result = execute_page( + sql, QueryExecutionOptions { - max_rows: Some(page_size), + max_rows: Some(if use_cursor { i32::MAX as usize } else { page_size }), + fetch_size: use_cursor.then_some(page_size), + page_size: use_cursor.then_some(page_size), + result_session_id: session_id.take(), + client_session_id: client_session_id.map(str::to_string), timeout_secs: options.timeout_secs, ..Default::default() }, ) .await?; + let fetched = result.rows.len(); + let complete = if use_cursor { + session_id = result.session_id.clone().filter(|session| !session.trim().is_empty()); + if result.truncated { + return Err("Incomplete cursor result during table export".to_string()); + } + if result.has_more && session_id.is_none() { + return Err("Result session ended before table export completed".to_string()); + } + !result.has_more + } else { + fetched < page_size + }; if !wrote_header { - write_csv_text_row(&mut writer, result.columns, options.csv_quote_mode)?; + write_csv_text_row(writer, result.columns, options.csv_quote_mode)?; wrote_header = true; } - let fetched = result.rows.len(); - if fetched == 0 { - break; - } for row in result.rows { writer.write_all(b"\n").map_err(|err| err.to_string())?; - write_csv_value_row(&mut writer, row, options.csv_quote_mode)?; + write_csv_value_row(writer, row, options.csv_quote_mode)?; } rows_exported += fetched as u64; - if fetched < page_size { + if complete { break; } offset += fetched; @@ -122,9 +175,168 @@ pub async fn export_table_data_csv_core(state: &AppState, options: TableCsvExpor #[cfg(test)] mod tests { - use super::{export_table_data_csv_core, CsvQuoteMode, TableCsvExportOptions}; + use super::{export_table_data_csv_core, write_table_csv_pages, CsvQuoteMode, TableCsvExportOptions}; use crate::connection::AppState; use crate::models::connection::{ConnectionConfig, DatabaseType}; + use crate::types::QueryResult; + + fn export_options(page_size: usize) -> TableCsvExportOptions { + TableCsvExportOptions { + file_path: String::new(), + connection_id: "csv-test".to_string(), + database: "app".to_string(), + schema: None, + table_name: "events".to_string(), + columns: vec!["id".to_string()], + page_size: Some(page_size), + timeout_secs: Some(30), + csv_quote_mode: CsvQuoteMode::default(), + } + } + + fn page(ids: &[usize], session_id: Option<&str>, has_more: bool) -> QueryResult { + serde_json::from_value(serde_json::json!({ + "columns": ["id"], + "rows": ids.iter().map(|id| vec![*id]).collect::>(), + "affected_rows": 0, + "execution_time_ms": 1, + "session_id": session_id, + "has_more": has_more + })) + .unwrap() + } + + #[tokio::test] + async fn cassandra_csv_export_follows_cursor_beyond_default_page_size() { + let first_ids: Vec = (1..=10_000).collect(); + let mut pages = vec![page(&first_ids, Some("cursor-1"), true), page(&[10_001], None, false)].into_iter(); + let mut requests = Vec::new(); + let mut output = Vec::new(); + let count = write_table_csv_pages( + &mut output, + DatabaseType::Cassandra, + &export_options(usize::MAX), + Some("export-test"), + |sql, options| { + requests.push((sql, options)); + std::future::ready(Ok(pages.next().expect("must not restart the first page"))) + }, + ) + .await + .unwrap(); + assert_eq!(count, 10_001); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].1.result_session_id, None); + assert_eq!(requests[1].1.result_session_id.as_deref(), Some("cursor-1")); + for (sql, options) in requests { + assert!(!sql.contains("OFFSET")); + assert!(!sql.contains("LIMIT")); + assert_eq!(options.max_rows, Some(i32::MAX as usize)); + assert_eq!(options.page_size, Some(10_000)); + assert_eq!(options.fetch_size, Some(10_000)); + assert_eq!(options.client_session_id.as_deref(), Some("export-test")); + assert_eq!(options.timeout_secs, Some(30)); + } + let csv = String::from_utf8(output).unwrap(); + assert_eq!(csv.lines().count(), 10_002); + assert_eq!(csv.lines().skip(1).collect::>().len(), 10_001); + } + + #[tokio::test] + async fn cassandra_csv_export_handles_short_empty_and_exact_cursor_pages() { + for pages in [ + vec![page(&[], None, false)], + vec![page(&[1, 2], None, false)], + vec![page(&[1], Some("cursor-1"), true), page(&[], Some("cursor-1"), true), page(&[2, 3], None, false)], + ] { + let expected_rows: usize = pages.iter().map(|page| page.rows.len()).sum(); + let expected_calls = pages.len(); + let mut pages = pages.into_iter(); + let mut calls = 0; + let count = write_table_csv_pages( + &mut Vec::new(), + DatabaseType::Cassandra, + &export_options(2), + Some("export-test"), + |_, _| { + calls += 1; + std::future::ready(Ok(pages.next().expect("unexpected extra request"))) + }, + ) + .await + .unwrap(); + assert_eq!(count, expected_rows as u64); + assert_eq!(calls, expected_calls); + } + } + + #[tokio::test] + async fn cassandra_csv_export_rejects_incomplete_cursor_pages() { + let mut truncated = page(&[1], None, false); + truncated.truncated = true; + for malformed in [page(&[1], None, true), page(&[1], Some(" "), true), truncated] { + let mut calls = 0; + let error = write_table_csv_pages( + &mut Vec::new(), + DatabaseType::Cassandra, + &export_options(2), + Some("export-test"), + |_, _| { + calls += 1; + std::future::ready(Ok(malformed.clone())) + }, + ) + .await + .expect_err("incomplete CSV must not be reported as successful"); + assert!(error.contains("export")); + assert_eq!(calls, 1); + } + } + + #[tokio::test] + async fn cassandra_csv_export_propagates_continuation_failure() { + let mut calls = 0; + let error = write_table_csv_pages( + &mut Vec::new(), + DatabaseType::Cassandra, + &export_options(2), + Some("export-test"), + |_, options| { + calls += 1; + std::future::ready(if calls == 1 { + Ok(page(&[1], Some("cursor-lost"), true)) + } else { + assert_eq!(options.result_session_id.as_deref(), Some("cursor-lost")); + Err("cursor lost".to_string()) + }) + }, + ) + .await + .unwrap_err(); + assert_eq!(error, "cursor lost"); + assert_eq!(calls, 2); + } + + #[tokio::test] + async fn non_cassandra_csv_export_preserves_limit_offset_paging() { + let mut pages = vec![page(&[1, 2], None, false), page(&[3], None, false)].into_iter(); + let mut sqls = Vec::new(); + let count = + write_table_csv_pages(&mut Vec::new(), DatabaseType::Sqlite, &export_options(2), None, |sql, options| { + sqls.push(sql); + assert_eq!(options.max_rows, Some(2)); + assert!(options.page_size.is_none()); + assert!(options.result_session_id.is_none()); + assert!(options.client_session_id.is_none()); + std::future::ready(Ok(pages.next().unwrap())) + }) + .await + .unwrap(); + assert_eq!(count, 3); + assert_eq!(sqls.len(), 2); + assert!(sqls[0].contains("LIMIT 2")); + assert!(sqls[1].contains("OFFSET 2")); + } fn salesforce_config(id: &str) -> ConnectionConfig { let mut config = serde_json::from_value::(serde_json::json!({ diff --git a/crates/dbx-sql-data/src/database_search_sql.rs b/crates/dbx-sql-data/src/database_search_sql.rs index d67ddf4870..a2a4d2aaad 100644 --- a/crates/dbx-sql-data/src/database_search_sql.rs +++ b/crates/dbx-sql-data/src/database_search_sql.rs @@ -338,6 +338,22 @@ mod tests { assert!(query.sql.contains("FROM \"public\".\"orders\"")); } + #[test] + fn cassandra_bounded_search_keeps_its_limit() { + let query = build_database_search_sql(DatabaseSearchSqlOptions { + database_type: Some(DatabaseType::Cassandra), + driver_profile: None, + schema: Some("app".to_string()), + table_name: "events".to_string(), + columns: vec![col("id", "int", true)], + term: "42".to_string(), + limit: Some(17), + }) + .unwrap(); + + assert!(query.sql.ends_with(" LIMIT 17;"), "sql was: {}", query.sql); + } + #[test] fn builds_oceanbase_oracle_search_query_with_rownum_limit() { let query = build_database_search_sql(DatabaseSearchSqlOptions { diff --git a/crates/dbx-sql-data/src/query_result_sql.rs b/crates/dbx-sql-data/src/query_result_sql.rs index 91107d193a..5dfd51eae1 100644 --- a/crates/dbx-sql-data/src/query_result_sql.rs +++ b/crates/dbx-sql-data/src/query_result_sql.rs @@ -335,6 +335,9 @@ pub fn build_count_query_sql(options: CountQuerySqlOptions) -> QuerySqlBuildResu if unsupported_pagination_type(options.database_type) { return err("unsupported"); } + if options.database_type == Some(DatabaseType::Cassandra) { + return err("unsupported"); + } let (execution_hint, statement) = split_leading_execution_hint(&statement, options.database_type); // A locking clause does not affect cardinality and cannot appear inside // every dialect's derived-table count query. PostgreSQL permits pagination @@ -4784,6 +4787,25 @@ WHERE u.id = picked.id; assert_eq!(plan.page_offset, Some(200)); } + #[test] + fn cassandra_offset_jump_keeps_agent_result_session_without_sql_pagination() { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM events".to_string(), + query_base_sql: "SELECT * FROM events".to_string(), + database_type: Some(DatabaseType::Cassandra), + pagination: QueryPagination { limit: 100, offset: 200, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert_eq!(plan.sql_to_execute, "SELECT * FROM events;"); + assert_eq!(plan.page_limit, Some(100)); + assert_eq!(plan.page_offset, Some(200)); + assert!(plan.page_sql.is_some()); + assert!(plan.count_sql.is_none()); + assert!(plan.use_agent_result_session); + } + #[test] fn rewritten_offset_jump_does_not_use_agent_result_session() { let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { diff --git a/crates/dbx-sql-dialect/src/sql_dialect/capabilities.rs b/crates/dbx-sql-dialect/src/sql_dialect/capabilities.rs index d0a26b23b8..dae88df56f 100644 --- a/crates/dbx-sql-dialect/src/sql_dialect/capabilities.rs +++ b/crates/dbx-sql-dialect/src/sql_dialect/capabilities.rs @@ -101,6 +101,11 @@ pub fn uses_single_row_insert_statements(database_type: DatabaseType) -> bool { pub fn pagination_strategy(database_type: Option, context: PaginationContext) -> TablePaginationStrategy { match database_type { Some(DatabaseType::Jdbc) => TablePaginationStrategy::AgentMaxRows, + Some(DatabaseType::Cassandra) + if matches!(context, PaginationContext::TablePreview | PaginationContext::UserQuery) => + { + TablePaginationStrategy::AgentMaxRows + } Some(DatabaseType::Oracle) if matches!(context, PaginationContext::TablePreview) => { TablePaginationStrategy::Rownum } diff --git a/crates/dbx-sql-dialect/src/sql_dialect/tests.rs b/crates/dbx-sql-dialect/src/sql_dialect/tests.rs index 1f6493d318..3e1ed2db21 100644 --- a/crates/dbx-sql-dialect/src/sql_dialect/tests.rs +++ b/crates/dbx-sql-dialect/src/sql_dialect/tests.rs @@ -199,6 +199,15 @@ fn maps_table_pagination_strategy_by_database_type() { assert_eq!(table_pagination_strategy(Some(DatabaseType::Questdb)), TablePaginationStrategy::QuestDbLimit); assert_eq!(table_pagination_strategy(Some(DatabaseType::Oracle)), TablePaginationStrategy::Rownum); assert_eq!(table_pagination_strategy(Some(DatabaseType::Oscar)), TablePaginationStrategy::Rownum); + assert_eq!(table_pagination_strategy(Some(DatabaseType::Cassandra)), TablePaginationStrategy::AgentMaxRows); + assert_eq!( + pagination_strategy(Some(DatabaseType::Cassandra), PaginationContext::BoundedRead), + TablePaginationStrategy::LimitOffset + ); + assert_eq!( + pagination_strategy(Some(DatabaseType::Cassandra), PaginationContext::UserQuery), + TablePaginationStrategy::AgentMaxRows + ); assert_eq!( pagination_strategy(Some(DatabaseType::Oracle), PaginationContext::BoundedRead), TablePaginationStrategy::Rownum @@ -221,6 +230,22 @@ fn maps_table_pagination_strategy_by_database_type() { assert_eq!(table_pagination_strategy(None), TablePaginationStrategy::LimitOffset); } +#[test] +fn cassandra_table_data_pagination_uses_the_agent_cursor() { + for offset in [0, 100, 200] { + assert_eq!( + build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Cassandra), + table_name: "paged_rows".to_string(), + limit: Some(100), + offset: Some(offset), + ..Default::default() + }), + "SELECT * FROM \"paged_rows\";" + ); + } +} + #[test] fn builds_select_sql_with_limit_syntax_for_database_type() { let columns = vec!["id".to_string(), "name".to_string()]; diff --git a/packages/app-tests/tableDataExport.test.ts b/packages/app-tests/tableDataExport.test.ts index 827ec2cb76..c6579b94e3 100644 --- a/packages/app-tests/tableDataExport.test.ts +++ b/packages/app-tests/tableDataExport.test.ts @@ -48,3 +48,139 @@ test("fetchTableDataForExport executes a VictoriaMetrics range query once", asyn assert.deepEqual(sqls, ['{__name__="flag"}[1h]']); assert.equal(exported.rows.length, 2); }); + +test("fetchTableDataForExport advances a Cassandra cursor past 10000 rows exactly once", async () => { + const firstRows = Array.from({ length: TABLE_DATA_EXPORT_PAGE_SIZE }, (_, index) => [index + 1]); + const requests: Array<{ sql: string; resultSessionId?: string }> = []; + const closed: Array = []; + const pages = [{ ...result(firstRows), session_id: "cursor-1", has_more: true }, result([[10_001], [10_002]])]; + + const exported = await fetchTableDataForExport({ + databaseType: "cassandra", + tableName: "events", + useAgentCursor: true, + buildPageSql: () => 'SELECT * FROM "events";', + executePage: async (sql, options) => { + requests.push({ sql, resultSessionId: options?.resultSessionId }); + return pages.shift() ?? result([]); + }, + closeCursor: async (sessionId) => { + closed.push(sessionId); + }, + }); + + assert.equal(exported.rows.length, TABLE_DATA_EXPORT_PAGE_SIZE + 2); + assert.equal(new Set(exported.rows.map((row) => row[0])).size, TABLE_DATA_EXPORT_PAGE_SIZE + 2); + assert.deepEqual(requests, [ + { sql: 'SELECT * FROM "events";', resultSessionId: undefined }, + { sql: 'SELECT * FROM "events";', resultSessionId: "cursor-1" }, + ]); + assert.deepEqual(closed, [undefined]); +}); + +test("fetchTableDataForExport trusts has_more for short and exact cursor pages", async () => { + const pages = [{ ...result([[1]]), session_id: "cursor-1", has_more: true }, { ...result([[2], [3]]), session_id: "cursor-1", has_more: true }, result([[4], [5]])]; + + const exported = await fetchTableDataForExport({ + databaseType: "cassandra", + tableName: "events", + pageSize: 2, + useAgentCursor: true, + buildPageSql: () => 'SELECT * FROM "events";', + executePage: async () => pages.shift() ?? result([]), + }); + + assert.deepEqual(exported.rows, [[1], [2], [3], [4], [5]]); +}); + +test("fetchTableDataForExport returns an empty terminal cursor page", async () => { + const exported = await fetchTableDataForExport({ + databaseType: "cassandra", + tableName: "events", + useAgentCursor: true, + buildPageSql: () => 'SELECT * FROM "events";', + executePage: async () => result([]), + }); + + assert.deepEqual(exported.rows, []); +}); + +test("fetchTableDataForExport rejects missing and malformed cursor sessions", async () => { + for (const sessionId of [undefined, " "]) { + let calls = 0; + await assert.rejects( + fetchTableDataForExport({ + databaseType: "cassandra", + tableName: "events", + useAgentCursor: true, + buildPageSql: () => 'SELECT * FROM "events";', + executePage: async () => { + calls += 1; + return { ...result([[1]]), session_id: sessionId, has_more: true }; + }, + }), + /session/i, + ); + assert.equal(calls, 1); + } +}); + +test("fetchTableDataForExport closes an active cursor after an error", async () => { + const closed: Array = []; + let calls = 0; + await assert.rejects( + fetchTableDataForExport({ + databaseType: "cassandra", + tableName: "events", + useAgentCursor: true, + buildPageSql: () => 'SELECT * FROM "events";', + executePage: async () => { + calls += 1; + if (calls === 1) return { ...result([[1]]), session_id: "cursor-1", has_more: true }; + throw new Error("cursor lost"); + }, + closeCursor: async (sessionId) => { + closed.push(sessionId); + }, + }), + /cursor lost/, + ); + + assert.deepEqual(closed, ["cursor-1"]); +}); + +test("fetchTableDataForExport refuses incomplete cursor pages and retains the active session for cleanup", async () => { + for (const malformed of [{ has_more: undefined }, { has_more: "false" }, { has_more: false, truncated: true }, { has_more: true, session_id: " " }]) { + const closed: Array = []; + let calls = 0; + await assert.rejects( + fetchTableDataForExport({ + databaseType: "cassandra", + tableName: "events", + useAgentCursor: true, + buildPageSql: () => 'SELECT * FROM "events";', + executePage: async () => { + calls += 1; + if (calls === 1) return { ...result([[1]]), session_id: "cursor-1", has_more: true }; + return { ...result([[2]]), ...malformed } as QueryResult; + }, + closeCursor: async (sessionId) => { + closed.push(sessionId); + }, + }), + /export/i, + ); + assert.equal(calls, 2); + assert.deepEqual(closed, ["cursor-1"]); + } +}); + +test("fetchTableDataForExport preserves non-cursor results without cursor metadata", async () => { + const exported = await fetchTableDataForExport({ + databaseType: "mysql", + tableName: "events", + buildPageSql: () => "SELECT id FROM events LIMIT 10000;", + executePage: async () => ({ ...result([[1]]), has_more: undefined }), + }); + assert.deepEqual(exported.rows, [[1]]); +});