fix(grid): guard editing and preserve transaction snapshots

This commit is contained in:
t8y2
2026-09-17 19:13:50 +08:00
parent 160ada25a1
commit 24ffa47bb5
11 changed files with 667 additions and 55 deletions
@@ -3,7 +3,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
import { useDataGridActions } from "@/composables/useDataGridActions";
import { clearTableMetadataCache } from "@/lib/metadata/tableMetadataCache";
import { restoredDataTabReloadFilters } from "@/lib/table/tableDataRefresh";
import type { QueryTab } from "@/types/database";
import type { IndexInfo, QueryTab } from "@/types/database";
const mocks = vi.hoisted(() => ({
buildTableSelectSql: vi.fn(),
@@ -79,14 +79,15 @@ vi.mock("@/stores/queryStore", () => ({
if (staleOrder) tab.orderByInput = undefined;
return true;
}),
setTableMeta: mocks.setTableMeta.mockImplementation((id: string, meta: NonNullable<QueryTab["tableMeta"]>) => {
setTableMeta: mocks.setTableMeta.mockImplementation((id: string, meta: NonNullable<QueryTab["tableMeta"]>, options: { rowIdentityPending?: boolean } = {}) => {
const tab = mocks.tabs.find((item) => item.id === id);
if (tab) {
tab.tableMeta = meta;
tab.tableMetaGeneration = mocks.metadataGeneration;
tab.tableMetaUpdatedAt = Date.now();
// 与真实 store 一致:仅真实元数据(columns 非空)落地才结束行标识等待
if (meta.columns.length > 0) tab.tableMetaPending = false;
if (options.rowIdentityPending) tab.tableMetaPending = true;
else if (meta.columns.length > 0) tab.tableMetaPending = false;
}
}),
}),
@@ -193,7 +194,7 @@ describe("useDataGridActions", () => {
releaseIndexes([]);
});
it("manual refresh reuses surviving primary keys intersected with fresh columns", async () => {
it("manual refresh derives primary keys from fresh columns", async () => {
const tab = tableDataTab();
mocks.tabs.push(tab);
// 表结构变化:旧主键 id 被删,新主键 id2
@@ -201,7 +202,102 @@ describe("useDataGridActions", () => {
const actions = useDataGridActions(computed(() => tab));
await actions.onReloadData(tab.id, tab.sql, "", "", "", 100, 0, "refresh");
// 旧 PK 不在新列中:等待段不得把失效 PK 写回 tab 元数据
expect(tab.tableMeta?.primaryKeys).toEqual(["id2"]);
});
it.each([true, false])("keeps a surviving old key read-only until index discovery settles (success=%s)", async (success) => {
const tab = reactive(tableDataTab());
mocks.tabs.push(tab);
mocks.getColumns.mockResolvedValue([{ ...tab.tableMeta!.columns[0], is_primary_key: false }]);
let resolveIndexes!: (indexes: IndexInfo[]) => void;
let rejectIndexes!: (error: Error) => void;
mocks.listIndexes.mockReturnValueOnce(
new Promise<IndexInfo[]>((resolve, reject) => {
resolveIndexes = resolve;
rejectIndexes = reject;
}),
);
const actions = useDataGridActions(computed(() => tab));
await actions.onReloadData(tab.id, tab.sql, "", "", "", 100, 0, "refresh");
expect(mocks.executeTabSql).toHaveBeenCalledOnce();
expect(tab.tableMetaPending).toBe(true);
expect(tab.tableMeta?.primaryKeys).toEqual([]);
if (success) {
resolveIndexes([]);
await vi.waitFor(() => expect(tab.tableMetaPending).toBe(false));
} else {
rejectIndexes(new Error("index discovery failed"));
await vi.waitFor(() => expect(mocks.setTableMeta).toHaveBeenCalledTimes(2));
expect(tab.tableMetaPending).toBe(true);
await actions.onReloadData(tab.id, tab.sql, "", "", "", 100, 0, "auto-refresh");
await vi.waitFor(() => expect(tab.tableMetaPending).toBe(false));
expect(mocks.listIndexes).toHaveBeenCalledTimes(2);
}
expect(tab.tableMeta?.primaryKeys).toEqual([]);
});
it("uses a freshly discovered unique index only after discovery finishes", async () => {
const tab = tableDataTab();
mocks.tabs.push(tab);
mocks.getColumns.mockResolvedValue([{ ...tab.tableMeta!.columns[0], is_primary_key: false }]);
let resolveIndexes!: (indexes: IndexInfo[]) => void;
mocks.listIndexes.mockReturnValueOnce(
new Promise<IndexInfo[]>((resolve) => {
resolveIndexes = resolve;
}),
);
const actions = useDataGridActions(computed(() => tab));
await actions.onReloadData(tab.id, tab.sql, "", "", "", 100, 0, "refresh");
expect(tab.tableMetaPending).toBe(true);
expect(tab.tableMeta?.primaryKeys).toEqual([]);
resolveIndexes([{ name: "users_id_unique", columns: ["id"], is_unique: true, is_primary: false }]);
await vi.waitFor(() => expect(tab.tableMetaPending).toBe(false));
expect(tab.tableMeta?.primaryKeys).toEqual(["id"]);
});
it.each(["TABLE", "VIEW", "MATERIALIZED_VIEW"])("preserves Oracle synthetic projection only for eligible tables (%s)", async (tableType) => {
mocks.getConfig.mockReturnValue({ id: "oracle-1", db_type: "oracle" });
const tab = tableDataTab({ connectionId: "oracle-1" });
tab.tableMeta!.tableType = tableType;
tab.tableMeta!.columns[0].is_primary_key = false;
tab.tableMeta!.primaryKeys = tableType === "TABLE" ? ["__DBX_ROWID"] : [];
mocks.tabs.push(tab);
mocks.getColumns.mockResolvedValue(tab.tableMeta!.columns);
let resolveIndexes!: (indexes: IndexInfo[]) => void;
mocks.listIndexes.mockReturnValueOnce(
new Promise<IndexInfo[]>((resolve) => {
resolveIndexes = resolve;
}),
);
const actions = useDataGridActions(computed(() => tab));
await actions.onReloadData(tab.id, tab.sql, "", "", "", 100, 0, "refresh");
expect(mocks.executeTabSql).toHaveBeenCalledOnce();
expect(mocks.buildTableSelectSql).toHaveBeenCalledWith(
expect.objectContaining({
includeRowId: tableType === "TABLE",
primaryKeys: tableType === "TABLE" ? ["__DBX_ROWID"] : [],
}),
);
expect(tab.tableMetaPending).toBe(true);
resolveIndexes([]);
await vi.waitFor(() => expect(tab.tableMetaPending).toBe(false));
expect(tab.tableMeta?.primaryKeys).toEqual(tableType === "TABLE" ? ["__DBX_ROWID"] : []);
});
it("keeps old metadata read-only when fresh columns fail", async () => {
const tab = tableDataTab();
mocks.tabs.push(tab);
mocks.getColumns.mockRejectedValueOnce(new Error("columns failed"));
const actions = useDataGridActions(computed(() => tab));
await expect(actions.onReloadData(tab.id, tab.sql, "", "", "", 100, 0, "refresh")).rejects.toThrow("columns failed");
expect(tab.tableMetaPending).toBe(true);
expect(mocks.executeTabSql).not.toHaveBeenCalled();
});
it("preserves rows when cache deletion fails and still executes", async () => {
@@ -50,6 +50,7 @@ function createEditorWithResult(
tableColumns?: Array<{ name: string; data_type: string; extra?: string; column_default?: string }>,
mongoCollectionGrid = false,
quickEntry = false,
editable = ref(true),
) {
let editor: ReturnType<typeof useDataGridEditor>;
const result = ref<{ columns: string[]; rows: CellValue[][] }>({
@@ -59,7 +60,7 @@ function createEditorWithResult(
editor = useDataGridEditor({
result: computed(() => result.value),
editable: computed(() => true),
editable: computed(() => editable.value),
databaseType: computed(() => (mongoCollectionGrid ? "mongodb" : "postgres")),
normalizeEditorInput: mongoCollectionGrid ? mongoDocumentGridInputValue : undefined,
connectionId: computed(() => "connection-1"),
@@ -145,6 +146,46 @@ describe("useDataGridEditor searched replacements", () => {
mocks.getConfig.mockReturnValue(undefined);
});
it.each([true, false])("blocks queued saves while row identity is pending (autoSave=%s)", async (autoSave) => {
const editable = ref(true);
const editor = createEditor(undefined, true, undefined, undefined, [["old", "keep", "last"]], undefined, undefined, false, false, editable);
editor.newRows.value = [];
editor.applyCellValue(0, 0, "new");
editable.value = false;
await editor.saveChanges({ autoSave });
expect(mocks.prepareDataGridSave).not.toHaveBeenCalled();
expect(mocks.executeBatch).not.toHaveBeenCalled();
expect(editor.hasPendingChanges.value).toBe(true);
editable.value = true;
mocks.prepareDataGridSave.mockResolvedValue({ statements: ["UPDATE people SET first='new'"], rollbackStatements: [] });
mocks.executeBatch.mockResolvedValue([]);
await editor.saveChanges();
expect(mocks.executeBatch).toHaveBeenCalledOnce();
});
it("rechecks row identity readiness after asynchronous save preparation", async () => {
const editable = ref(true);
const editor = createEditor(undefined, true, undefined, undefined, [["old", "keep", "last"]], undefined, undefined, false, false, editable);
editor.newRows.value = [];
editor.applyCellValue(0, 0, "new");
let finishPreparation!: (value: { statements: string[]; rollbackStatements: string[] }) => void;
mocks.prepareDataGridSave.mockReturnValueOnce(
new Promise((resolve) => {
finishPreparation = resolve;
}),
);
const save = editor.saveChanges();
expect(mocks.prepareDataGridSave).toHaveBeenCalledOnce();
editable.value = false;
finishPreparation({ statements: ["UPDATE people SET first='new'"], rollbackStatements: [] });
await save;
expect(mocks.executeBatch).not.toHaveBeenCalled();
expect(editor.hasPendingChanges.value).toBe(true);
expect(editor.isSaving.value).toBe(false);
});
it("stages one undoable batch, keeps empty strings and preserves prior pending edits", () => {
const editor = createEditor(undefined, true, undefined, undefined, [["hit", "keep", "hit"]]);
editor.newRows.value = [];
@@ -8,7 +8,7 @@ import { buildTableSelectSql, quoteTableDataIdentifier } from "@/lib/table/table
import { tableOpenPageLimit } from "@/lib/table/tableOpenPageLimit";
import { tableDataLargeValuePreviewOptions } from "@/lib/dataGrid/dataGridLargeValues";
import { elasticsearchCursorPageJumpRequestCount } from "@/lib/dataGrid/dataGridPagination";
import { shouldIncludeSyntheticRowId } from "@/lib/table/tableEditing";
import { editablePrimaryKeys, shouldIncludeSyntheticRowId } from "@/lib/table/tableEditing";
import { tableMetaForDataTab } from "@/lib/table/tableDataTabMeta";
import * as api from "@/lib/backend/api";
import type { ColumnInfo, QueryTab } from "@/types/database";
@@ -138,10 +138,6 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
});
}
// columnsOnly:只重拉列(getColumns),跳过 listIndexes。供手动刷新的等待段
// 使用——列投影必须新鲜才能构建 SELECT,而主键/索引用旧值交集构建本次查询
// 是安全的(PK 名只进保护集合,不进 SQL 文本),完整索引元数据由随后的
// 后台全量刷新补齐(列命中 columns facet 缓存,只花一次 listIndexes)。
async function refreshDataTabTableMeta(tab: QueryTab, options: { force?: boolean; columnsOnly?: boolean; trace?: { traceId: string; elapsed: () => string } } = {}): Promise<boolean> {
if (tab.mode !== "data" || !tab.connectionId || !tab.database) return false;
const tableMeta = tableMetaForDataTab(tab);
@@ -159,6 +155,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
// still receive its result. Only the latest request may update this tab.
const requestToken = {};
metadataRequests.set(tab, requestToken);
tab.tableMetaPending = true;
const metadataGenerationAtStart = connectionStore.metadataGenerationFor(target.connectionId, target.database);
const trace = options.trace;
@@ -175,7 +172,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
// 复用共享表元数据缓存(30s TTL + in-flight 去重),多个入口对同一张表
// 不再各自往返 getColumns/listIndexes。跨连接生命周期的强制重建走 force,
// 避免同一共享缓存把断链前的旧列再次交回本次 reload。
const loaded: { columns: TableMetadataColumns; primaryKeys: string[] } = options.columnsOnly
const loaded: { columns: TableMetadataColumns; primaryKeys: string[]; rowIdentityResolved: boolean } = options.columnsOnly
? await (async () => {
const { columns } = await loadTableColumns({
connectionId: target.connectionId,
@@ -188,9 +185,8 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
catalog: target.catalog,
force: options.force === true,
});
const available = new Set(columns.map((column) => column.name.toLocaleLowerCase()));
const primaryKeys = (tab.tableMeta?.primaryKeys ?? []).filter((key) => available.has(key.toLocaleLowerCase()));
return { columns, primaryKeys };
const primaryKeys = editablePrimaryKeys(effectiveDatabaseTypeForConnection(config), columns, target.tableType);
return { columns, primaryKeys, rowIdentityResolved: false };
})()
: await (async () => {
const { metadata } = await loadTableMetadata({
@@ -204,7 +200,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
catalog: target.catalog,
force: options.force === true,
});
return { columns: metadata.columns, primaryKeys: metadata.primaryKeys };
return { columns: metadata.columns, primaryKeys: metadata.primaryKeys, rowIdentityResolved: metadata.rowIdentityResolved !== false };
})();
const columns = loaded.columns;
console.info("[DBX][reloadData:metadata:get-columns:done]", { traceId: trace?.traceId, elapsed: trace?.elapsed(), columnCount: columns.length });
@@ -229,7 +225,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
return false;
}
const primaryKeys = loaded.primaryKeys;
queryStore.setTableMeta(target.tabId, {
const refreshedMeta = {
catalog: target.catalog,
database: target.database,
schema: target.schema,
@@ -237,7 +233,12 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
tableType: target.tableType,
columns,
primaryKeys,
});
};
if (loaded.rowIdentityResolved) {
queryStore.setTableMeta(target.tabId, refreshedMeta);
} else {
queryStore.setTableMeta(target.tabId, refreshedMeta, { rowIdentityPending: true });
}
return true;
}
@@ -319,7 +320,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
}
const connectionGeneration = connectionStore.metadataGenerationFor(tab.connectionId, tab.database);
const lifecycleStale = isDataTabMetadataLifecycleStale(tab, connectionGeneration);
const shouldRefreshMetadata = lifecycleStale || !hasRealTableMetaColumns || metadataAgeMs > DATA_TAB_METADATA_TTL_MS;
const shouldRefreshMetadata = lifecycleStale || tab.tableMetaPending || !hasRealTableMetaColumns || metadataAgeMs > DATA_TAB_METADATA_TTL_MS;
// Dameng 元数据必须与数据查询串行(同 useSidebarDataOpenRuntime),
// 延后到查询完成后再启动。主动刷新和跨生命周期重建除外:必须先拿到新列再
// 构建 SQL,否则第一次 toolbar reload 仍会沿用断链前的显式列列表。
@@ -336,7 +337,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
});
};
if (lifecycleStale || intent === "refresh") {
if (!hasRealTableMetaColumns) tab.tableMetaPending = true;
tab.tableMetaPending = true;
console.info("[DBX][reloadData:metadata:await:start]", { traceId, elapsed: elapsed(), reason: intent === "refresh" ? "manual-refresh" : "lifecycle-stale", metadataAgeMs });
try {
if (intent === "refresh") {
@@ -381,8 +382,6 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
incomingSortMissing = rebuiltColumnNames.length > 0 && simpleDataGridOrderByReferencesMissingColumn(orderBy, rebuiltColumnNames);
if (incomingSortMissing) tab.orderByInput = undefined;
if (intent === "refresh" && !lifecycleStale) {
// 后台补齐索引/主键元数据:列 facet 刚写入缓存,这里只花一次
// listIndexes 往返;完成后 PK/可编辑性恢复最新(此前用旧 PK 交集)。
void refreshDataTabTableMeta(tab, { force: false, trace: { traceId, elapsed } })
.then(() => {
console.info("[DBX][reloadData:metadata:background-indexes:done]", { traceId, elapsed: elapsed() });
@@ -1849,6 +1849,7 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
}
async function saveChanges(saveOptions: SaveChangesOptions = {}) {
if (!editable.value) return;
if (saveOptions.autoSave && manualSaveRequired.value) return;
if (isSaving.value) {
if (saveOptions.autoSave) pendingAutoSaveRequested = true;
@@ -1881,6 +1882,7 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
});
if (!confirmed) return;
}
if (!editable.value) return;
if (customHandler && snapshot.newRows.length > 0 && customHandler.supportsInsert !== true && customHandler.canInsert !== true) {
saveError.value = i18n.global.t("grid.insertRowsNotSupported");
return;
@@ -1977,6 +1979,10 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) {
return;
}
}
if (!editable.value || stmtOptions?.tableMeta !== tableMeta.value) {
await finishInterruptedSaveChanges(snapshot);
return;
}
const start = Date.now();
let apiResult: { affected_rows?: number } | undefined;
console.info("[DBX][dataGrid:save-statements]", {
@@ -55,6 +55,20 @@ describe("tableMetadataCache columns facet request counts", () => {
mocks.listIndexes.mockResolvedValue([]);
});
it("does not treat failed index discovery as a complete cacheable row identity", async () => {
mocks.getColumns.mockResolvedValue([{ ...column("id"), is_primary_key: false }]);
mocks.listIndexes.mockRejectedValueOnce(new Error("indexes unavailable"));
const failed = await loadTableMetadata(usersRequest);
expect(failed.metadata.rowIdentityResolved).toBe(false);
expect(failed.metadata.primaryKeys).toEqual([]);
expect(getCachedTableMetadata(usersRequest)).toBeUndefined();
const retried = await loadTableMetadata(usersRequest);
expect(retried.metadata.rowIdentityResolved).toBe(true);
expect(mocks.listIndexes).toHaveBeenCalledTimes(2);
expect(mocks.getColumns).toHaveBeenCalledOnce();
});
it("R1 — cold single-source columns load: 1 getColumns, 0 listIndexes", async () => {
const result = await loadTableColumns({ ...usersRequest });
expect(result.columns[0]?.name).toBe("users");
@@ -17,6 +17,7 @@ export interface TableMetadata {
columns: ColumnInfo[];
indexes: IndexInfo[];
primaryKeys: string[];
rowIdentityResolved?: boolean;
cachedAt: number;
}
@@ -331,11 +332,16 @@ export async function loadTableMetadata(request: TableMetadataRequest): Promise<
// discovery independently unless an agent-backed PostgreSQL-family
// relation must first report its visible schema for the index lookup.
const resolveReportedSchema = (request.databaseType === "vastbase" || request.databaseType === "kingbase") && !request.schema;
const indexesPromise = resolveReportedSchema ? undefined : loadTableIndexes(request).catch((): IndexInfo[] => []);
let rowIdentityResolved = true;
const indexFailure = (): IndexInfo[] => {
rowIdentityResolved = false;
return [];
};
const indexesPromise = resolveReportedSchema ? undefined : loadTableIndexes(request).catch(indexFailure);
const columnsResult = await columnsPromise;
const columns = columnsResult.columns;
const resolvedSchema = resolveReportedSchema ? columns.find((column) => column.resolved_schema)?.resolved_schema : request.schema;
const indexes = columns.length > 0 ? await (indexesPromise ?? loadTableIndexes({ ...request, schema: resolvedSchema }).catch((): IndexInfo[] => [])) : [];
const indexes = columns.length > 0 ? await (indexesPromise ?? loadTableIndexes({ ...request, schema: resolvedSchema }).catch(indexFailure)) : [];
const primaryKeys = editableRowIdentifierColumns(request.databaseType as DatabaseType, columns, indexes, request.tableType);
return {
schema: resolvedSchema || undefined,
@@ -346,6 +352,7 @@ export async function loadTableMetadata(request: TableMetadataRequest): Promise<
columns,
indexes,
primaryKeys,
rowIdentityResolved: rowIdentityResolved && columns.length > 0,
cachedAt: columnsResult.cachedAt,
};
},
@@ -353,7 +360,7 @@ export async function loadTableMetadata(request: TableMetadataRequest): Promise<
);
// 必须在 unregister 前比较:最后一个在途加载注销时会顺带清掉代数记录
if (invalidationStampAtStart === (tableMetadataInvalidationStamps.get(scopeKey) ?? 0)) {
if (metadata.rowIdentityResolved && invalidationStampAtStart === (tableMetadataInvalidationStamps.get(scopeKey) ?? 0)) {
tableMetadataCache.set(scope, metadata, { cachedAt: metadata.cachedAt });
}
} finally {
@@ -17,6 +17,25 @@ describe("data tab snapshot restore vs tableMetaPending", () => {
setActivePinia(createPinia());
});
it("keeps projection-only metadata pending until complete identity metadata arrives", async () => {
const { useQueryStore } = await import("@/stores/queryStore");
const store = useQueryStore();
const tabId = store.createTab("pg-1", "app", "users", "data", "public");
const tab = store.tabs.find((item) => item.id === tabId)!;
const metadata = { tableName: "users", schema: "public", columns: [column("id")], primaryKeys: ["id"] };
store.setTableMeta(tabId, metadata);
expect(tab.tableMetaPending).toBe(false);
store.setTableMeta(tabId, { ...metadata, columns: [{ ...column("id"), is_primary_key: false }], primaryKeys: [] }, { rowIdentityPending: true });
expect(tab.tableMetaPending).toBe(true);
expect(tab.tableMeta?.primaryKeys).toEqual([]);
store.setTableMeta(tabId, { ...metadata, columns: [], primaryKeys: [] });
expect(tab.tableMetaPending).toBe(true);
store.setTableMeta(tabId, metadata);
expect(tab.tableMetaPending).toBe(false);
});
it("does not roll back real metadata to a placeholder snapshot and re-pends placeholder restores", async () => {
vi.doMock("@/lib/tabs/tabResultCache", async (importOriginal) => {
const actual = await importOriginal<typeof import("@/lib/tabs/tabResultCache")>();
+3 -2
View File
@@ -5064,9 +5064,10 @@ export const useQueryStore = defineStore("query", () => {
return clearInvalidDataTabSortState(tab, tab.tableMeta.columns);
}
function setTableMeta(id: string, meta: NonNullable<QueryTab["tableMeta"]>) {
function setTableMeta(id: string, meta: NonNullable<QueryTab["tableMeta"]>, options: { rowIdentityPending?: boolean } = {}) {
const tab = tabs.value.find((t) => t.id === id);
if (tab) {
if (options.rowIdentityPending) tab.tableMetaPending = true;
tab.tableMeta = meta;
// 记录写入时的连接元数据代次:disconnect/关库/死池重连会使该代次递增,
// 代次失配视同冷缓存,即使位于 30s TTL 窗口内也会重建结构(issue #6623 /
@@ -5078,7 +5079,7 @@ export const useQueryStore = defineStore("query", () => {
// columns/primaryKeys 为空的占位身份(如 useNavigationTargets),不得
// 借此提前解除编辑门控。失败/中止路径不清除——标签页保持只读是安全
// 兜底,刷新或重开表会重新加载元数据恢复
if (meta.columns.length > 0) tab.tableMetaPending = false;
if (meta.columns.length > 0 && !options.rowIdentityPending) tab.tableMetaPending = false;
}
}
+3 -14
View File
@@ -269,26 +269,15 @@ pub struct TransactionSession {
pub pool_key: String,
pub last_activity: std::time::Instant,
pub busy: bool,
pub snapshot_rotation_safe: bool,
pub connection_id: String,
pub database: String,
pub schema: Option<String>,
}
impl TransactionSession {
/// Whether this session holds a natively connected MySQL/PostgreSQL
/// connection whose transaction keeps a read snapshot that a rollback +
/// fresh BEGIN can rotate in place. Agent and external-driver sessions are
/// excluded: their rollback path closes the dedicated session, so they
/// cannot reopen in place.
pub fn connection_is_native_snapshot_dialect(&self) -> bool {
// The connection is inspected under its own lock elsewhere; here only
// the variant is needed, and reading the enum discriminant without the
// lock would race. Callers hold the sessions lock and pass the same
// decision point, so mirror them by locking briefly.
match self.connection.try_lock() {
Ok(conn) => matches!(&*conn, TxnConnection::Mysql(Some(_)) | TxnConnection::Postgres(_)),
Err(_) => false,
}
pub fn can_rotate_read_only_snapshot(&self, conn: &TxnConnection) -> bool {
self.snapshot_rotation_safe && matches!(conn, TxnConnection::Mysql(Some(_)) | TxnConnection::Postgres(_))
}
}
+258 -15
View File
@@ -5613,6 +5613,7 @@ async fn begin_transaction_session(
pool_key: pool_key.clone(),
last_activity: std::time::Instant::now(),
busy: false,
snapshot_rotation_safe: !consistent_snapshot,
connection_id: connection_id.to_string(),
database: database.to_string(),
schema: schema.map(|s| s.to_string()),
@@ -5786,6 +5787,9 @@ pub async fn execute_in_manual_transaction_with_options(
// session remains intact.
check_read_only_for_connection_multi(state, &pool_key, &statements).await?;
let classification: Vec<bool> =
classify_manual_transaction_statements(db_type, statements.len(), options.classification_sql.as_deref());
let connection = {
let mut sessions = state.transaction_sessions.write().await;
let Some(session) = sessions.get_mut(txn_session_id) else {
@@ -5799,6 +5803,8 @@ pub async fn execute_in_manual_transaction_with_options(
Some(session)
} else {
session.busy = true;
session.snapshot_rotation_safe &=
classification.len() == statements.len() && classification.iter().all(|proven| *proven);
session.last_activity = std::time::Instant::now();
None
}
@@ -5820,15 +5826,6 @@ pub async fn execute_in_manual_transaction_with_options(
let row_limit = options.max_rows.unwrap_or(MAX_ROWS).max(1);
let mut results = Vec::with_capacity(statements.len());
// Sticky-dialect classification pairing. The core splits both the execution
// SQL and, when present, the user-facing classification SQL with the same
// dialect-aware splitter. A marker is emitted only when both lists have the
// same non-zero count and every paired user statement is proven read-only;
// any mismatch is fail-closed (no marker). This is deliberately a
// trust-boundary count/position pairing, not a SQL-equivalence parser.
let classification: Vec<bool> =
classify_manual_transaction_statements(db_type, statements.len(), options.classification_sql.as_deref());
let mut conn = connection.lock().await;
for (i, statement) in statements.iter().enumerate() {
let result = match &mut *conn {
@@ -5897,7 +5894,7 @@ pub async fn execute_in_manual_transaction_with_options(
// a fresh session) instead of failing the already-successful batch.
if manual_txn_batch_fully_proven_read_only(&results) && {
let sessions = state.transaction_sessions.read().await;
sessions.get(txn_session_id).map(|session| session.connection_is_native_snapshot_dialect()).unwrap_or(false)
sessions.get(txn_session_id).map(|session| session.can_rotate_read_only_snapshot(&conn)).unwrap_or(false)
} {
if let Err(rotation_error) = rotate_clean_read_only_snapshot(&mut conn, schema).await {
let removed = {
@@ -5971,6 +5968,7 @@ where
Some(sessions.remove(txn_session_id).expect("session exists").connection)
} else {
session.busy = true;
session.snapshot_rotation_safe = false;
session.last_activity = std::time::Instant::now();
None
}
@@ -6123,9 +6121,8 @@ where
stream_result
}
/// Whether a finished manual-transaction batch was fully proven read-only, so
/// the transaction still holds only a read view and no user data state. Factored
/// out so the snapshot-rotation decision is unit-testable without a database.
/// Whether a finished batch was fully proven read-only. Session history is
/// checked separately before rotating its snapshot.
fn manual_txn_batch_fully_proven_read_only(results: &[ExecuteMultiResult]) -> bool {
!results.is_empty() && results.iter().all(|result| result.manual_transaction_proven_read_only)
}
@@ -6137,8 +6134,8 @@ fn manual_txn_batch_fully_proven_read_only(results: &[ExecuteMultiResult]) -> bo
/// read-only session would otherwise never see rows committed by others, and the
/// clean-state toolbar hides Commit/Rollback, leaving disconnect/reconnect as
/// the only visible way out. A rollback of a read-only transaction and a fresh
/// BEGIN are both cheap metadata operations, and the batch gate keeps this off
/// every write path. Failures here never fail the (already successful) batch:
/// BEGIN are both cheap metadata operations. The session-history gate excludes
/// any prior write or uncertain batch. Failures never fail the successful batch:
/// the session is torn down instead and the frontend's existing
/// rolled-back-session recovery transparently opens a new one on the next run.
async fn rotate_clean_read_only_snapshot(conn: &mut TxnConnection, schema: Option<&str>) -> Result<(), String> {
@@ -6579,6 +6576,250 @@ mod tests {
use super::*;
use crate::query_cancel::RunningTaskMetadata;
mod manual_transaction_snapshot_tests {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
struct SnapshotTestConnect {
commands: Arc<Mutex<Vec<String>>>,
fail_once: Option<&'static str>,
}
impl deadpool_postgres::Connect for SnapshotTestConnect {
fn connect(
&self,
config: &tokio_postgres::Config,
) -> futures::future::BoxFuture<
'_,
Result<(tokio_postgres::Client, tokio::task::JoinHandle<()>), tokio_postgres::Error>,
> {
let mut config = config.clone();
config.user("test").ssl_mode(tokio_postgres::config::SslMode::Disable);
let commands = Arc::clone(&self.commands);
let mut fail_once = self.fail_once;
Box::pin(async move {
let (client_socket, mut server_socket) = tokio::io::duplex(8192);
tokio::spawn(async move {
let startup_length = server_socket.read_u32().await.unwrap();
let mut startup = vec![0; startup_length as usize - 4];
server_socket.read_exact(&mut startup).await.unwrap();
server_socket.write_all(b"R\0\0\0\x08\0\0\0\0Z\0\0\0\x05I").await.unwrap();
let mut statement = String::new();
while let Ok(message_type) = server_socket.read_u8().await {
let length = server_socket.read_u32().await.unwrap();
let mut payload = vec![0; length as usize - 4];
server_socket.read_exact(&mut payload).await.unwrap();
let (response_type, response) = match message_type {
b'P' => {
statement =
String::from_utf8(payload.split(|byte| *byte == 0).nth(1).unwrap().to_vec())
.unwrap();
(b'1', Vec::new())
}
b'B' => (b'2', Vec::new()),
b'D' => {
if payload.first() == Some(&b'S') {
server_socket.write_all(b"t\0\0\0\x06\0\0").await.unwrap();
}
(b'n', Vec::new())
}
b'E' => {
commands.lock().unwrap().push(statement.clone());
if fail_once == Some(statement.as_str()) {
fail_once = None;
(b'E', b"SERROR\0CXX000\0Mtest statement failed\0\0".to_vec())
} else {
let tag = if statement.starts_with("UPDATE") {
"UPDATE 1"
} else if statement.starts_with("SELECT") {
"SELECT 0"
} else {
statement.as_str()
};
(b'C', format!("{tag}\0").into_bytes())
}
}
b'S' => (b'Z', vec![b'I']),
b'C' => (b'3', Vec::new()),
b'X' => break,
other => panic!("unexpected PostgreSQL message: {other}"),
};
server_socket.write_u8(response_type).await.unwrap();
server_socket.write_u32((response.len() + 4) as u32).await.unwrap();
server_socket.write_all(&response).await.unwrap();
}
});
let (client, connection) = config.connect_raw(client_socket, tokio_postgres::NoTls).await?;
let task = tokio::spawn(async move {
let _ = connection.await;
});
Ok((client, task))
})
}
}
async fn snapshot_state(
fail_once: Option<&'static str>,
) -> (AppState, Arc<Mutex<Vec<String>>>, tempfile::TempDir) {
let directory = tempfile::tempdir().unwrap();
let state = AppState::new(Storage::open(&directory.path().join("storage.db")).await.unwrap());
let config = test_connection_config(DatabaseType::Postgres);
state.configs.write().await.insert(config.id.clone(), config);
let commands = Arc::new(Mutex::new(Vec::new()));
let manager = deadpool_postgres::Manager::from_connect(
tokio_postgres::Config::new(),
SnapshotTestConnect { commands: Arc::clone(&commands), fail_once },
deadpool_postgres::ManagerConfig::default(),
);
let pool = deadpool_postgres::Pool::builder(manager).max_size(1).build().unwrap();
let connection = timeout(Duration::from_secs(5), pool.get()).await.unwrap().unwrap();
state.transaction_sessions.write().await.insert(
"snapshot".to_string(),
TransactionSession {
connection: Arc::new(tokio::sync::Mutex::new(TxnConnection::Postgres(Box::new(connection)))),
pool_key: "conn-1".to_string(),
last_activity: std::time::Instant::now(),
busy: false,
snapshot_rotation_safe: true,
connection_id: "conn-1".to_string(),
database: "test".to_string(),
schema: None,
},
);
(state, commands, directory)
}
async fn execute_snapshot_batch(
state: &AppState,
sql: &str,
classification_sql: Option<&str>,
) -> Result<Vec<ExecuteMultiResult>, String> {
timeout(
Duration::from_secs(5),
execute_in_manual_transaction_with_options(
state,
"snapshot",
sql,
"test",
None,
ManualTransactionExecutionOptions {
classification_sql: classification_sql.map(str::to_string),
..Default::default()
},
),
)
.await
.expect("manual transaction must not deadlock")
}
#[tokio::test]
async fn snapshot_predicate_uses_the_already_held_connection_guard() {
let (state, _, _directory) = snapshot_state(None).await;
let connection = Arc::clone(&state.transaction_sessions.read().await["snapshot"].connection);
let guard = connection.lock().await;
assert!(connection.try_lock().is_err());
let mut sessions = state.transaction_sessions.write().await;
let session = sessions.get_mut("snapshot").unwrap();
assert!(session.can_rotate_read_only_snapshot(&guard));
session.snapshot_rotation_safe = false;
assert!(!session.can_rotate_read_only_snapshot(&guard));
session.snapshot_rotation_safe = true;
assert!(!session.can_rotate_read_only_snapshot(&TxnConnection::Mysql(None)));
}
#[tokio::test]
async fn repeated_reads_rotate_and_no_statement_preserves_clean_history() {
let (state, commands, _directory) = snapshot_state(None).await;
execute_snapshot_batch(&state, "-- no statement", None).await.unwrap();
assert!(commands.lock().unwrap().is_empty());
for _ in 0..2 {
let results = execute_snapshot_batch(&state, "SELECT 1", Some("SELECT 1")).await.unwrap();
assert!(manual_txn_batch_fully_proven_read_only(&results));
let sessions = state.transaction_sessions.read().await;
assert!(!sessions["snapshot"].busy);
assert!(sessions["snapshot"].snapshot_rotation_safe);
}
assert_eq!(*commands.lock().unwrap(), ["SELECT 1", "ROLLBACK", "BEGIN", "SELECT 1", "ROLLBACK", "BEGIN"]);
}
#[tokio::test]
async fn writes_unknown_and_mixed_batches_permanently_prevent_rotation() {
for (sql, classification) in [
("UPDATE users SET id = 2", Some("UPDATE users SET id = 2")),
("SELECT 1", None),
("SELECT 1", Some("SELECT unknown_function()")),
("SELECT 1; UPDATE users SET id = 2", Some("SELECT 1; UPDATE users SET id = 2")),
] {
let (state, commands, _directory) = snapshot_state(None).await;
execute_snapshot_batch(&state, sql, classification).await.unwrap();
for _ in 0..2 {
execute_snapshot_batch(&state, "SELECT 1", Some("SELECT 1")).await.unwrap();
}
{
let sessions = state.transaction_sessions.read().await;
assert!(!sessions["snapshot"].snapshot_rotation_safe);
assert!(!sessions["snapshot"].busy);
}
assert!(!commands.lock().unwrap().iter().any(|sql| sql == "ROLLBACK" || sql == "BEGIN"));
commit_manual_transaction(&state, "snapshot").await.unwrap();
assert_eq!(commands.lock().unwrap().last().unwrap(), "COMMIT");
}
}
#[tokio::test]
async fn unclassified_stream_prevents_later_snapshot_rotation() {
let (state, commands, _directory) = snapshot_state(None).await;
stream_rows_in_manual_transaction(&state, "snapshot", "SELECT 1", 10, |_| Ok(())).await.unwrap();
execute_snapshot_batch(&state, "SELECT 1", Some("SELECT 1")).await.unwrap();
assert!(!state.transaction_sessions.read().await["snapshot"].snapshot_rotation_safe);
assert_eq!(*commands.lock().unwrap(), ["SELECT 1", "SELECT 1"]);
}
#[tokio::test]
async fn rotation_failure_removes_session_without_failing_successful_read() {
for failure in ["ROLLBACK", "BEGIN"] {
let (state, commands, _directory) = snapshot_state(Some(failure)).await;
let results = execute_snapshot_batch(&state, "SELECT 1", Some("SELECT 1")).await.unwrap();
assert!(manual_txn_batch_fully_proven_read_only(&results));
assert!(!state.transaction_sessions.read().await.contains_key("snapshot"));
let expected = if failure == "BEGIN" {
vec!["SELECT 1", "ROLLBACK", "BEGIN", "ROLLBACK"]
} else {
vec!["SELECT 1", "ROLLBACK", "ROLLBACK"]
};
assert_eq!(*commands.lock().unwrap(), expected);
}
}
#[tokio::test]
async fn failed_batch_after_write_keeps_existing_rollback_cleanup() {
let (state, commands, _directory) = snapshot_state(Some("SELECT 1")).await;
execute_snapshot_batch(&state, "UPDATE users SET id = 2", Some("UPDATE users SET id = 2")).await.unwrap();
let error = execute_snapshot_batch(&state, "SELECT 1", Some("SELECT 1")).await.unwrap_err();
assert!(error.contains("manual transaction was rolled back"));
assert!(!state.transaction_sessions.read().await.contains_key("snapshot"));
assert_eq!(*commands.lock().unwrap(), ["UPDATE users SET id = 2", "SELECT 1", "ROLLBACK"]);
}
#[tokio::test]
async fn cancelled_stream_after_write_keeps_existing_rollback_cleanup() {
let (state, commands, _directory) = snapshot_state(None).await;
execute_snapshot_batch(&state, "UPDATE users SET id = 2", Some("UPDATE users SET id = 2")).await.unwrap();
let cancel = CancellationToken::new();
cancel.cancel();
let result =
stream_rows_in_manual_transaction_with_cancel(&state, "snapshot", "SELECT 1", 10, Some(cancel), |_| {
Ok(())
})
.await;
assert!(result.is_err());
assert!(!state.transaction_sessions.read().await.contains_key("snapshot"));
let commands = commands.lock().unwrap();
assert_eq!(commands.last().unwrap(), "ROLLBACK");
assert!(!commands.iter().any(|sql| sql == "BEGIN"));
}
}
#[test]
fn manual_txn_batch_proven_read_only_requires_non_empty_all_proven_results() {
// Empty batch (e.g. comments-only script) must not rotate: there is no
@@ -9583,6 +9824,7 @@ for line in sys.stdin:
pool_key: pool_key.to_string(),
last_activity: std::time::Instant::now(),
busy: false,
snapshot_rotation_safe: true,
connection_id: "jdbc-conn".to_string(),
database: "dbx_test".to_string(),
schema: None,
@@ -10778,6 +11020,7 @@ for line in sys.stdin:
pool_key: "agent-conn".to_string(),
last_activity: std::time::Instant::now(),
busy: false,
snapshot_rotation_safe: true,
connection_id: "agent-conn".to_string(),
database: "ORCL".to_string(),
schema: None,
@@ -0,0 +1,197 @@
use dbx_core::connection::AppState;
use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
use dbx_core::query::{
begin_database_backup_snapshot, begin_manual_transaction, commit_manual_transaction, execute_in_manual_transaction,
execute_in_manual_transaction_with_options, execute_sql_statement, rollback_manual_transaction,
ManualTransactionExecutionOptions,
};
use dbx_core::storage::Storage;
use std::sync::Arc;
fn live_config(prefix: &str, database_type: DatabaseType, default_port: u16) -> ConnectionConfig {
serde_json::from_value(serde_json::json!({
"id": format!("snapshot-rotation-{}", uuid::Uuid::new_v4().simple()),
"name": "Snapshot rotation regression",
"db_type": database_type,
"host": std::env::var(format!("{prefix}_HOST")).expect("live database host"),
"port": std::env::var(format!("{prefix}_PORT")).ok().and_then(|port| port.parse::<u16>().ok()).unwrap_or(default_port),
"username": std::env::var(format!("{prefix}_USER")).expect("live database user"),
"password": std::env::var(format!("{prefix}_PASSWORD")).expect("live database password"),
"database": std::env::var(format!("{prefix}_DATABASE")).expect("live database name"),
"connect_timeout_secs": 5,
"query_timeout_secs": 30,
"idle_timeout_secs": 60,
"keepalive_interval_secs": 0,
"url_params": std::env::var(format!("{prefix}_URL_PARAMS")).ok()
}))
.expect("live connection configuration")
}
async fn setup(config: ConnectionConfig) -> (Arc<AppState>, std::path::PathBuf, String) {
let storage_path = std::env::temp_dir().join(format!("dbx-snapshot-rotation-{}.db", uuid::Uuid::new_v4().simple()));
let state = Arc::new(AppState::new(Storage::open(&storage_path).await.expect("temporary storage")));
let table_name = format!("dbx_pr9434_{}", uuid::Uuid::new_v4().simple());
let database = config.database.clone().expect("database");
state.configs.write().await.insert(config.id.clone(), config.clone());
execute_sql_statement(
&state,
&config.id,
&database,
&format!("CREATE TABLE {table_name} (id INTEGER PRIMARY KEY)"),
None,
None,
)
.await
.expect("create isolated regression table");
execute_sql_statement(
&state,
&config.id,
&database,
&format!("INSERT INTO {table_name} (id) VALUES (1)"),
None,
None,
)
.await
.expect("seed regression table");
(state, storage_path, table_name)
}
async fn assert_transaction_row_count(
state: &AppState,
transaction: &str,
database: &str,
table_name: &str,
expected: i64,
) {
let sql = format!("SELECT COUNT(*) AS row_count FROM {table_name}");
let result = execute_in_manual_transaction_with_options(
state,
transaction,
&sql,
database,
None,
ManualTransactionExecutionOptions {
max_rows: Some(10),
classification_sql: Some(sql.clone()),
..Default::default()
},
)
.await
.expect("read transaction rows");
assert!(result[0].manual_transaction_proven_read_only);
assert_row_count(&result[0].result.rows, expected);
}
fn assert_row_count(rows: &[Vec<serde_json::Value>], expected: i64) {
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].len(), 1);
let count = rows[0][0].as_i64().or_else(|| rows[0][0].as_str().and_then(|value| value.parse().ok()));
assert_eq!(count, Some(expected));
}
async fn cleanup(state: &AppState, config: &ConnectionConfig, table_name: &str, storage_path: std::path::PathBuf) {
execute_sql_statement(
state,
&config.id,
config.database.as_deref().expect("database"),
&format!("DROP TABLE {table_name}"),
None,
None,
)
.await
.expect("drop isolated regression table");
let _ = std::fs::remove_file(storage_path);
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_MANUAL_TXN_MYSQL_* on the remote test server"]
async fn live_mysql_clean_manual_reads_refresh_snapshot_after_external_write() {
let config = live_config("DBX_LIVE_MANUAL_TXN_MYSQL", DatabaseType::Mysql, 3306);
let database = config.database.clone().expect("database");
let (state, storage_path, table_name) = setup(config.clone()).await;
let transaction =
begin_manual_transaction(&state, &config.id, &database, None, None).await.expect("begin transaction");
assert_transaction_row_count(&state, &transaction, &database, &table_name, 1).await;
execute_sql_statement(
&state,
&config.id,
&database,
&format!("INSERT INTO {table_name} (id) VALUES (2)"),
None,
None,
)
.await
.expect("external committed write");
assert_transaction_row_count(&state, &transaction, &database, &table_name, 2).await;
rollback_manual_transaction(&state, &transaction).await.expect("rollback read-only transaction");
cleanup(&state, &config, &table_name, storage_path).await;
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_MANUAL_TXN_MYSQL_* on the remote test server"]
async fn live_mysql_explicit_backup_snapshot_stays_repeatable() {
let config = live_config("DBX_LIVE_MANUAL_TXN_MYSQL", DatabaseType::Mysql, 3306);
let database = config.database.clone().expect("database");
let (state, storage_path, table_name) = setup(config.clone()).await;
let transaction =
begin_database_backup_snapshot(&state, &config.id, &database).await.expect("begin consistent snapshot");
assert_transaction_row_count(&state, &transaction, &database, &table_name, 1).await;
execute_sql_statement(
&state,
&config.id,
&database,
&format!("INSERT INTO {table_name} (id) VALUES (2)"),
None,
None,
)
.await
.expect("external committed write");
assert_transaction_row_count(&state, &transaction, &database, &table_name, 1).await;
rollback_manual_transaction(&state, &transaction).await.expect("rollback backup snapshot");
cleanup(&state, &config, &table_name, storage_path).await;
}
async fn verify_write_then_read_survives_commit(config: ConnectionConfig) {
let database = config.database.clone().expect("database");
let (state, storage_path, table_name) = setup(config.clone()).await;
let transaction =
begin_manual_transaction(&state, &config.id, &database, None, None).await.expect("begin transaction");
execute_in_manual_transaction(
&state,
&transaction,
&format!("INSERT INTO {table_name} (id) VALUES (2)"),
&database,
None,
Some(10),
)
.await
.expect("uncommitted write");
assert_transaction_row_count(&state, &transaction, &database, &table_name, 2).await;
assert_transaction_row_count(&state, &transaction, &database, &table_name, 2).await;
commit_manual_transaction(&state, &transaction).await.expect("commit must retain prior write");
let result = execute_sql_statement(
&state,
&config.id,
&database,
&format!("SELECT COUNT(*) AS row_count FROM {table_name}"),
None,
None,
)
.await
.expect("read committed rows");
assert_row_count(&result.rows, 2);
cleanup(&state, &config, &table_name, storage_path).await;
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_MANUAL_TXN_MYSQL_* on the remote test server"]
async fn live_mysql_write_then_reads_preserve_commit() {
verify_write_then_read_survives_commit(live_config("DBX_LIVE_MANUAL_TXN_MYSQL", DatabaseType::Mysql, 3306)).await;
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_MANUAL_TXN_POSTGRES_* on the remote test server"]
async fn live_postgres_write_then_reads_preserve_commit() {
verify_write_then_read_survives_commit(live_config("DBX_LIVE_MANUAL_TXN_POSTGRES", DatabaseType::Postgres, 5432))
.await;
}