diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index d22df9bd4..448b05bd4 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -394,6 +394,19 @@ function currentTableStructureDatabaseType(): DatabaseType | undefined { return props.node.connectionId ? tableStructureDatabaseTypeForConnection(connectionStore.getConfig(props.node.connectionId)) : undefined; } +// Database types whose backend size query is implemented (schema.rs +// `database_size_core`): the native MySQL pool (incl. JDBC/agent MySQL) and the +// native PostgreSQL pool family. Agent-backed PG forks (HighGo/Vastbase/ +// Kingbase) are NOT dispatched by the backend, so the manual action must not be +// offered for them — otherwise the menu appears but can never return a size. +const DATABASE_SIZE_SUPPORTED_TYPES = new Set(["mysql", "postgres", "gaussdb", "kwdb", "opengauss"]); + +const canFetchDatabaseSize = computed(() => { + if (props.node.type !== "database" || props.node.catalog) return false; + const dbType = currentDatabaseType(); + return !!dbType && DATABASE_SIZE_SUPPORTED_TYPES.has(dbType); +}); + function rawDatabaseType(): DatabaseType | undefined { return props.node.connectionId ? connectionStore.getConfig(props.node.connectionId)?.db_type : undefined; } @@ -1782,6 +1795,16 @@ async function clearNodeDefaultDatabase() { } } +async function fetchDatabaseSize() { + const node = props.node; + if (!node.connectionId || !node.database) return; + try { + await connectionStore.fetchDatabaseSize(node.connectionId, node.database); + } catch (e: any) { + toast(t("contextMenu.fetchDatabaseSizeFailed", { message: e?.message || String(e) }), 5000); + } +} + async function refresh() { try { await connectionStore.refreshTreeNode(props.node); @@ -4227,6 +4250,14 @@ const tableComment = computed(() => : null, ); const labelWidthClass = computed(() => treeLabelWidthClass({ fullWidth: usesFullWidthLabel.value, hasTrailingComment: !!columnComment.value || !!tableComment.value })); +const formattedSize = computed(() => { + const bytes = props.node.sizeBytes; + if (bytes == null || bytes < 0) return null; + if (bytes < 1024) return `${bytes} B`; + if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)} KB`; + if (bytes < 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; + return `${(bytes / (1024 * 1024 * 1024)).toFixed(2)} GB`; +}); const paddingLeft = computed(() => treeItemPaddingLeft(props.depth)); const tableSearchParentId = computed(() => props.node.tableSearchParentId || ""); const tableSearchValue = computed(() => { @@ -5227,6 +5258,9 @@ function treeItemMenuItems(): ContextMenuItem[] { items.push({ label: t("contextMenu.clearDefaultDatabase"), action: clearNodeDefaultDatabase, icon: Database }); } } + if (canFetchDatabaseSize.value) { + items.push({ label: t("contextMenu.fetchDatabaseSize"), action: fetchDatabaseSize, icon: Database }); + } if (canEditDatabaseProperties.value) { items.push({ label: t("contextMenu.editDatabaseProperties"), action: openEditDatabasePropertiesDialog, icon: SquarePen }); } @@ -5751,6 +5785,7 @@ function treeItemMenuItems(): ContextMenuItem[] { {{ columnComment }} {{ tableComment }} + {{ formattedSize }} {{ t("connection.readOnlyBadge") }} diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 66599f915..f6fc8392a 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -1450,6 +1450,8 @@ export default { finalProxyPortCopied: "Final proxy port {port} copied", setDefaultDatabase: "Set as Default Database", clearDefaultDatabase: "Clear Default Database", + fetchDatabaseSize: "Fetch Database Size", + fetchDatabaseSizeFailed: "Failed to fetch database size: {message}", exportDatabase: "Export Database", exportAllDatabases: "Export All Databases", exportDatabaseSuccess: "Exported {count} tables, up to {limit} rows each", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index 8016ebeab..245b2b224 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -1392,6 +1392,8 @@ export default withEnglishFallback({ finalProxyPortCopied: "Puerto proxy final {port} copiado", setDefaultDatabase: "Establecer como base de datos predeterminada", clearDefaultDatabase: "Limpiar base de datos predeterminada", + fetchDatabaseSize: "Obtener tamaño de la base de datos", + fetchDatabaseSizeFailed: "Error al obtener el tamaño de la base de datos: {message}", exportDatabase: "Exportar base de datos", exportAllDatabases: "Exportar todas las bases de datos", exportDatabaseSuccess: "Se exportaron {count} tablas, hasta {limit} filas cada una", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index 9bd146992..24fcfaaff 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -1390,6 +1390,8 @@ export default withEnglishFallback({ finalProxyPortCopied: "Porta proxy finale {port} copiata", setDefaultDatabase: "Imposta come Database Predefinito", clearDefaultDatabase: "Cancella Database Predefinito", + fetchDatabaseSize: "Ottieni Dimensione Database", + fetchDatabaseSizeFailed: "Impossibile ottenere la dimensione del database: {message}", exportDatabase: "Esporta Database", exportAllDatabases: "Esporta tutti i database", exportDatabaseSuccess: "Esportate {count} tabelle, fino a {limit} righe ciascuna", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 63fa79c17..f71084f09 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -1389,6 +1389,8 @@ export default withEnglishFallback({ finalProxyPortCopied: "最終プロキシポート {port} をコピーしました", setDefaultDatabase: "デフォルトデータベースに設定", clearDefaultDatabase: "デフォルトデータベースをクリア", + fetchDatabaseSize: "データベースサイズを取得", + fetchDatabaseSizeFailed: "データベースサイズの取得に失敗しました：{message}", exportDatabase: "データベースをエクスポート", exportAllDatabases: "すべてのデータベースをエクスポート", exportDatabaseSuccess: "{count}テーブルをエクスポートしました（各{limit}行まで）", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index b2c7b6db3..b804275b7 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -1392,6 +1392,8 @@ export default withEnglishFallback({ copyName: "Copiar Nome", setDefaultDatabase: "Definir como Banco de Dados Padrão", clearDefaultDatabase: "Limpar Banco de Dados Padrão", + fetchDatabaseSize: "Obter Tamanho do Banco de Dados", + fetchDatabaseSizeFailed: "Falha ao obter o tamanho do banco de dados: {message}", exportDatabase: "Exportar Banco de Dados", exportAllDatabases: "Exportar todos os bancos de dados", exportDatabaseSuccess: "{count} tabelas exportadas, até {limit} linhas cada", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index a58c23f26..ae8b8d5bd 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -1452,6 +1452,8 @@ export default withEnglishFallback({ finalProxyPortCopied: "最终代理端口 {port} 已复制", setDefaultDatabase: "设为默认数据库", clearDefaultDatabase: "取消默认数据库", + fetchDatabaseSize: "获取数据库大小", + fetchDatabaseSizeFailed: "获取数据库大小失败：{message}", exportDatabase: "导出数据库", exportAllDatabases: "导出全部数据库", exportDatabaseSuccess: "已导出 {count} 张表，每表最多 {limit} 行", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index d657c0642..f7c1b6bd9 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -1392,6 +1392,8 @@ export default withEnglishFallback({ finalProxyPortCopied: "最終代理連接埠 {port} 已複製", setDefaultDatabase: "設為預設資料庫", clearDefaultDatabase: "取消預設資料庫", + fetchDatabaseSize: "取得資料庫大小", + fetchDatabaseSizeFailed: "取得資料庫大小失敗：{message}", exportDatabase: "匯出資料庫", exportAllDatabases: "匯出全部資料庫", exportDatabaseSuccess: "已匯出 {count} 張資料表，每張最多 {limit} 行", diff --git a/apps/desktop/src/lib/__tests__/databaseTree.spec.ts b/apps/desktop/src/lib/__tests__/databaseTree.spec.ts new file mode 100644 index 000000000..26df71b27 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/databaseTree.spec.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "vitest"; +import { buildDatabaseTreeNodes, sortSidebarDatabases } from "@/lib/database/databaseTree"; + +describe("buildDatabaseTreeNodes", () => { + it("passes size_bytes from DatabaseInfo to tree node sizeBytes", () => { + const databases = [ + { name: "app_db", size_bytes: 1048576 }, + { name: "metrics", size_bytes: 5368709120 }, + ]; + const nodes = buildDatabaseTreeNodes("conn-1", databases); + expect(nodes).toHaveLength(2); + expect(nodes.find((n) => n.label === "app_db")?.sizeBytes).toBe(1048576); + expect(nodes.find((n) => n.label === "metrics")?.sizeBytes).toBe(5368709120); + }); + + it("sets sizeBytes to null when size_bytes is undefined (drivers without size)", () => { + const databases = [{ name: "my_db" }]; + const nodes = buildDatabaseTreeNodes("conn-1", databases); + expect(nodes).toHaveLength(1); + expect(nodes[0].sizeBytes).toBeNull(); + }); + + it("sets sizeBytes to null when size_bytes is null (no CONNECT privilege)", () => { + const databases = [{ name: "restricted_db", size_bytes: null }]; + const nodes = buildDatabaseTreeNodes("conn-1", databases as any); + expect(nodes).toHaveLength(1); + expect(nodes[0].sizeBytes).toBeNull(); + }); + + it("still builds nodes when size_bytes is missing for all databases", () => { + const databases = [{ name: "db1" }, { name: "db2" }, { name: "db3" }]; + const nodes = buildDatabaseTreeNodes("conn-1", databases); + expect(nodes).toHaveLength(3); + expect(nodes.every((n) => n.sizeBytes === null)).toBe(true); + }); +}); + +describe("sortSidebarDatabases", () => { + it("sorts databases alphabetically with numeric awareness", () => { + const databases = [{ name: "db10" }, { name: "db2" }, { name: "db1" }]; + const sorted = sortSidebarDatabases(databases); + expect(sorted.map((d) => d.name)).toEqual(["db1", "db2", "db10"]); + }); +}); diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index e3f3dfeb8..d4637c0d6 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -127,6 +127,8 @@ export const syncSavedSqlDirectory = forward("syncSavedSqlDirectory"); // Schema export const listDatabases = forward("listDatabases"); +export const listDatabaseStatistics = forward("listDatabaseStatistics"); +export const databaseSize = forward("databaseSize"); export const listDorisCatalogs = forward("listDorisCatalogs"); export const listDorisCatalogDatabases = forward("listDorisCatalogDatabases"); export const listSqlServerLinkedServers = forward("listSqlServerLinkedServers"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index a0a3216b2..5260eb32d 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -1,6 +1,7 @@ import type { ConnectionConfig, DatabaseInfo, + DatabaseStatistics, SchemaInfo, LinkedServerInfo, CatalogInfo, @@ -521,6 +522,14 @@ export async function listDatabases(connectionId: string): Promise { + return get(`/api/schema/database-statistics?${qs({ connection_id: connectionId })}`); +} + +export async function databaseSize(connectionId: string, database: string): Promise { + return get(`/api/schema/database-size?${qs({ connection_id: connectionId, database })}`); +} + export async function listDorisCatalogs(connectionId: string): Promise { return get(`/api/schema/doris/catalogs?${qs({ connection_id: connectionId })}`); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 57297e2bf..e822c8df7 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -3,6 +3,7 @@ import { listen, type UnlistenFn } from "@tauri-apps/api/event"; import type { ConnectionConfig, DatabaseInfo, + DatabaseStatistics, SchemaInfo, LinkedServerInfo, CatalogInfo, @@ -665,6 +666,14 @@ export async function listDatabases(connectionId: string): Promise { + return invoke("list_database_statistics", { connectionId }); +} + +export async function databaseSize(connectionId: string, database: string): Promise { + return invoke("database_size", { connectionId, database }); +} + export async function listDorisCatalogs(connectionId: string): Promise { return invoke("list_doris_catalogs", { connectionId }); } diff --git a/apps/desktop/src/lib/database/databaseTree.ts b/apps/desktop/src/lib/database/databaseTree.ts index 937e0deca..6f7d3566c 100644 --- a/apps/desktop/src/lib/database/databaseTree.ts +++ b/apps/desktop/src/lib/database/databaseTree.ts @@ -26,6 +26,7 @@ export function buildDatabaseTreeNodes(connectionId: string, databases: Database type: "database" as const, connectionId, database: name, + sizeBytes: db.size_bytes ?? null, isExpanded: false, children: [], }, diff --git a/apps/desktop/src/stores/__tests__/connectionStore.databaseSize.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.databaseSize.spec.ts new file mode 100644 index 000000000..b5a905ebe --- /dev/null +++ b/apps/desktop/src/stores/__tests__/connectionStore.databaseSize.spec.ts @@ -0,0 +1,116 @@ +import { createPinia, setActivePinia } from "pinia"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ConnectionConfig, TreeNode } from "@/types/database"; + +function installLocalStorage() { + const data = new Map(); + vi.stubGlobal("localStorage", { + getItem: vi.fn((key: string) => data.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => data.set(key, value)), + removeItem: vi.fn((key: string) => data.delete(key)), + }); +} + +function mysqlConnection(): ConnectionConfig { + return { + id: "mysql-1", + name: "MySQL", + db_type: "mysql", + host: "127.0.0.1", + port: 3306, + username: "root", + password: "", + database: "app", + } as ConnectionConfig; +} + +function treeWithDatabases(connection: ConnectionConfig, databases: { name: string; sizeBytes: number | null }[]): TreeNode[] { + return [ + { + id: connection.id, + label: connection.name, + type: "connection", + connectionId: connection.id, + isExpanded: true, + children: databases.map((db) => ({ + id: `${connection.id}:${db.name}`, + label: db.name, + type: "database", + connectionId: connection.id, + database: db.name, + sizeBytes: db.sizeBytes, + isExpanded: false, + children: [], + })), + }, + ]; +} + +async function setupStore(databaseSizeImpl: (connectionId: string, database: string) => Promise) { + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + databaseSize: vi.fn(databaseSizeImpl), + listDatabaseStatistics: vi.fn().mockResolvedValue([]), + saveSchemaCache: vi.fn().mockResolvedValue(undefined), + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + loadSchemaCache: vi.fn().mockResolvedValue(null), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); + const { useConnectionStore } = await import("@/stores/connectionStore"); + return useConnectionStore(); +} + +describe("connectionStore database size (manual)", () => { + beforeEach(() => { + vi.resetModules(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + }); + + it("fetchDatabaseSize updates only the targeted database node", async () => { + const store = await setupStore(async () => 2048); + const connection = mysqlConnection(); + store.connections = [connection]; + store.connectedIds.add(connection.id); + store.treeNodes = treeWithDatabases(connection, [ + { name: "app", sizeBytes: null }, + { name: "other", sizeBytes: 111 }, + ]); + + const result = await store.fetchDatabaseSize(connection.id, "app"); + + expect(result).toBe(2048); + const children = store.treeNodes[0].children!; + expect(children.find((n) => n.database === "app")?.sizeBytes).toBe(2048); + expect(children.find((n) => n.database === "other")?.sizeBytes).toBe(111); + }); + + it("automatic statistics with null does not overwrite a manually fetched size", async () => { + const store = await setupStore(async () => 4096); + const connection = mysqlConnection(); + store.connections = [connection]; + store.connectedIds.add(connection.id); + store.treeNodes = treeWithDatabases(connection, [{ name: "app", sizeBytes: 4096 }]); + + // Automatic pass timed out for "app" (size_bytes null) — must keep the manual value. + const changed = store.applyDatabaseStatisticsToConnectionTree(connection.id, [{ name: "app", size_bytes: null }]); + + expect(changed).toBe(false); + expect(store.treeNodes[0].children![0].sizeBytes).toBe(4096); + }); + + it("automatic statistics still fills sizes that are not yet known", async () => { + const store = await setupStore(async () => 4096); + const connection = mysqlConnection(); + store.connections = [connection]; + store.connectedIds.add(connection.id); + store.treeNodes = treeWithDatabases(connection, [{ name: "app", sizeBytes: null }]); + + const changed = store.applyDatabaseStatisticsToConnectionTree(connection.id, [{ name: "app", size_bytes: 9000 }]); + + expect(changed).toBe(true); + expect(store.treeNodes[0].children![0].sizeBytes).toBe(9000); + }); +}); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 97ffa5984..73758bd7e 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -1,7 +1,23 @@ import { defineStore } from "pinia"; import { uuid } from "@/lib/common/utils"; import { ref, computed, watch, markRaw } from "vue"; -import type { ColumnInfo, CompletionAssistantCandidate, CompletionAssistantObjectKind, CompletionAssistantRequest, ConnectionConfig, CatalogInfo, ForeignKeyInfo, ObjectInfo, SchemaInfo, SidebarLayout, TableInfo, TreeNode, TunnelProfile, VectorCollectionMeta } from "@/types/database"; +import type { + ColumnInfo, + CompletionAssistantCandidate, + CompletionAssistantObjectKind, + CompletionAssistantRequest, + ConnectionConfig, + CatalogInfo, + DatabaseStatistics, + ForeignKeyInfo, + ObjectInfo, + SchemaInfo, + SidebarLayout, + TableInfo, + TreeNode, + TunnelProfile, + VectorCollectionMeta, +} from "@/types/database"; import { applyPinnedTreeNodeState, inheritNaturalTreeNodeOrder, migrateLegacyPinnedTreeNodeIds, syncPinnedTreeNodeStateInPlace, treeNodePinKey } from "@/lib/app/pinnedItems"; import { reconcileLayout, @@ -85,6 +101,12 @@ const DISCONNECT_REQUEST_TIMEOUT_MS = 5_000; const DEFAULT_KEEPALIVE_INTERVAL_SECS = 30; const METADATA_LIST_PAGE_CACHE_TTL_MS = 30_000; const METADATA_LIST_PAGE_CACHE_MAX_ENTRIES = 160; +const DATABASE_STATISTICS_CACHE_TTL_MS = 5 * 60_000; +const DATABASE_STATISTICS_CACHE_MAX_ENTRIES = 64; +const DATABASE_STATISTICS_TIMEOUT_MS = 10_000; +// Manual, single-database size fetch (right-click). One schema is much cheaper +// than the automatic all-database pass, so it gets a far more generous budget. +const MANUAL_DATABASE_SIZE_TIMEOUT_MS = 60_000; export const COMPLETION_METADATA_CONCURRENCY = 2; const MONGO_LEGACY_DRIVER_PROFILE = "mongodb-legacy"; const MONGO_LEGACY_DRIVER_LABEL = "MongoDB (Legacy)"; @@ -319,6 +341,10 @@ export const useConnectionStore = defineStore("connection", () => { ttlMs: METADATA_LIST_PAGE_CACHE_TTL_MS, maxEntries: METADATA_LIST_PAGE_CACHE_MAX_ENTRIES, }); + const databaseStatisticsCache = new MetadataResultCache({ + ttlMs: DATABASE_STATISTICS_CACHE_TTL_MS, + maxEntries: DATABASE_STATISTICS_CACHE_MAX_ENTRIES, + }); const metadataTraceLogger: MetadataLoadTraceLogger = (event) => { console.debug("[DBX][metadata-load:trace]", event); }; @@ -696,6 +722,38 @@ export const useConnectionStore = defineStore("connection", () => { } } + async function withDatabaseStatisticsTimeout(promise: Promise): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => { + reject(new Error(`Connection timed out while loading database statistics after ${Math.ceil(DATABASE_STATISTICS_TIMEOUT_MS / 1000)}s.`)); + }, DATABASE_STATISTICS_TIMEOUT_MS); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } + } + + async function withDatabaseSizeTimeout(promise: Promise): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => { + reject(new Error(`Timed out while fetching database size after ${Math.ceil(MANUAL_DATABASE_SIZE_TIMEOUT_MS / 1000)}s.`)); + }, MANUAL_DATABASE_SIZE_TIMEOUT_MS); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } + } + async function withDisconnectRequestTimeout(connectionId: string, promise: Promise): Promise { let timedOut = false; let timer: ReturnType | undefined; @@ -1087,8 +1145,120 @@ export const useConnectionStore = defineStore("connection", () => { }; } + function supportsAsyncDatabaseStatistics(config?: ConnectionConfig): boolean { + // Must match the backend's database-size dispatch (schema.rs): native MySQL + // pool (incl. JDBC/agent MySQL) and the native PostgreSQL pool family. Agent + // -backed PG forks (HighGo/Vastbase/Kingbase) are NOT dispatched, so they are + // intentionally excluded — advertising them would show sizes the backend + // cannot compute. + const dbType = effectiveDatabaseTypeForConnection(config); + return dbType === "mysql" || dbType === "postgres" || dbType === "gaussdb" || dbType === "kwdb" || dbType === "opengauss"; + } + + function databaseStatisticsCacheScope(connectionId: string): MetadataScopeInput { + return { + kind: "database-statistics", + connectionId, + driverProfile: metadataDriverProfile(getConfig(connectionId)), + }; + } + + async function loadDatabaseStatisticsCached(connectionId: string, options?: { force?: boolean }): Promise { + const scope = databaseStatisticsCacheScope(connectionId); + if (!options?.force) { + const cached = databaseStatisticsCache.get(scope); + if (cached) return cached.value; + } + return runTreeMetadataLoad( + scope, + async () => { + if (!connectedIds.value.has(connectionId)) return []; + if (!options?.force) { + const cached = databaseStatisticsCache.get(scope); + if (cached) return cached.value; + } + console.debug("[DBX][database-statistics:start]", { connectionId, force: options?.force === true }); + const startedAt = Date.now(); + const statistics = await withDatabaseStatisticsTimeout(api.listDatabaseStatistics(connectionId)); + databaseStatisticsCache.set(scope, statistics); + console.debug("[DBX][database-statistics:done]", { + connectionId, + resultCount: statistics.length, + elapsedMs: Date.now() - startedAt, + }); + return statistics; + }, + // A forced tree refresh should bypass the cached value, but it must not + // start a second expensive statistics query while one is already active. + { ...options, force: false }, + ); + } + + function applyDatabaseStatisticsToConnectionTree(connectionId: string, statistics: readonly DatabaseStatistics[]): boolean { + const connectionNode = findNode(treeNodes.value, connectionId); + if (!connectionNode?.children) return false; + + const sizeByDatabase = new Map(statistics.map((entry) => [entry.name, entry.size_bytes ?? null] as const)); + let changed = false; + for (const child of connectionNode.children) { + if (child.type !== "database" || child.catalog || child.database == null || !sizeByDatabase.has(child.database)) continue; + const nextSize = sizeByDatabase.get(child.database) ?? null; + // Don't let an automatic pass that timed out (null) erase a size the user + // fetched manually via the right-click action. Only apply null when the + // node has no size yet. + if (nextSize === null && child.sizeBytes != null) continue; + if (child.sizeBytes !== nextSize) { + child.sizeBytes = nextSize; + changed = true; + } + } + return changed; + } + + async function fetchDatabaseSize(connectionId: string, database: string): Promise { + const cacheKey = schemaCacheKey(connectionId, "databases"); + const size = await withDatabaseSizeTimeout(api.databaseSize(connectionId, database)); + const connectionNode = findNode(treeNodes.value, connectionId); + const child = connectionNode?.children?.find((node) => node.type === "database" && !node.catalog && node.database === database); + if (child && child.sizeBytes !== (size ?? null)) { + child.sizeBytes = size ?? null; + if (connectionNode?.children) { + void savePersistedConnectionTreeChildren(cacheKey, connectionNode.children); + } + } + return size ?? null; + } + + function scheduleDatabaseStatisticsLoad(connectionId: string, cacheKey: string, options?: { force?: boolean }) { + if (!supportsAsyncDatabaseStatistics(getConfig(connectionId))) return; + window.setTimeout(() => { + if (!connectedIds.value.has(connectionId)) return; + void loadDatabaseStatisticsCached(connectionId, options) + .then((statistics) => { + if (statistics.length === 0) return; + if (!applyDatabaseStatisticsToConnectionTree(connectionId, statistics)) return; + const node = findNode(treeNodes.value, connectionId); + if (node?.children) { + void savePersistedConnectionTreeChildren(cacheKey, node.children); + } + }) + .catch((error) => { + console.warn("[DBX][database-statistics:load-failed]", { connectionId, error }); + }); + }, 0); + } + function invalidateMetadataCaches(match: MetadataCacheInvalidation): number { - return metadataListPageCache.invalidate(match) + invalidateTableMetadataCache(match); + let removed = metadataListPageCache.invalidate(match) + invalidateTableMetadataCache(match); + if (match.connectionId) { + removed += databaseStatisticsCache.invalidate({ + kind: "database-statistics", + connectionId: match.connectionId, + }); + } else { + removed += databaseStatisticsCache.invalidate(match); + } + return removed; } function invalidateMetadataCachesByTreePrefix(prefix: string) { @@ -2253,6 +2423,7 @@ export const useConnectionStore = defineStore("connection", () => { if (!options?.force) { const cached = await loadPersistedTreeChildren(node, cacheKey); if (cached.hit) { + scheduleDatabaseStatisticsLoad(connectionId, cacheKey); if (cached.isStale) refreshStaleTreeNode(node); return; } @@ -2292,6 +2463,7 @@ export const useConnectionStore = defineStore("connection", () => { if (!canApplyTreeMetadataResult(node)) return; setChildren(node, children); await savePersistedConnectionTreeChildren(cacheKey, node.children || children); + scheduleDatabaseStatisticsLoad(connectionId, cacheKey, { force: options?.force }); } } node.isExpanded = true; @@ -5281,6 +5453,8 @@ export const useConnectionStore = defineStore("connection", () => { setDefaultDatabase, clearDefaultDatabase, isDefaultDatabase, + fetchDatabaseSize, + applyDatabaseStatisticsToConnectionTree, setVisibleDatabases, clearVisibleDatabases, ensureVisibleDatabase, diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 876f754f9..55da9b999 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -324,6 +324,12 @@ export interface JdbcPluginStatus { export interface DatabaseInfo { name: string; + size_bytes?: number | null; +} + +export interface DatabaseStatistics { + name: string; + size_bytes: number | null; } export interface SchemaInfo { @@ -667,6 +673,7 @@ export interface TreeNode { signature?: string; tableType?: string; comment?: string | null; + sizeBytes?: number | null; objectCount?: number; loadedKeyCount?: number; totalKeyCount?: number; diff --git a/crates/dbx-core/src/db/clickhouse_driver.rs b/crates/dbx-core/src/db/clickhouse_driver.rs index 583f600f8..73be1f827 100644 --- a/crates/dbx-core/src/db/clickhouse_driver.rs +++ b/crates/dbx-core/src/db/clickhouse_driver.rs @@ -478,7 +478,11 @@ pub async fn test_connection(client: &ChClient, timeout: Duration) -> Result<(), pub async fn list_databases(client: &ChClient) -> Result, String> { let result = ch_query(client, "SELECT name FROM system.databases ORDER BY name", None).await?; - Ok(result.data.iter().map(|row| DatabaseInfo { name: row[0].as_str().unwrap_or("").to_string() }).collect()) + Ok(result + .data + .iter() + .map(|row| DatabaseInfo { name: row[0].as_str().unwrap_or("").to_string(), size_bytes: None }) + .collect()) } pub async fn list_tables(client: &ChClient, database: &str) -> Result, String> { diff --git a/crates/dbx-core/src/db/cloudflare_d1/mod.rs b/crates/dbx-core/src/db/cloudflare_d1/mod.rs index d6080073e..c97b60229 100644 --- a/crates/dbx-core/src/db/cloudflare_d1/mod.rs +++ b/crates/dbx-core/src/db/cloudflare_d1/mod.rs @@ -111,7 +111,7 @@ pub async fn connect(config: &ConnectionConfig, timeout: Duration) -> Result Result, String> { - Ok(vec![DatabaseInfo { name: "main".to_string() }]) + Ok(vec![DatabaseInfo { name: "main".to_string(), size_bytes: None }]) } pub async fn list_tables(client: &CloudflareD1Client, _schema: &str) -> Result, String> { diff --git a/crates/dbx-core/src/db/influxdb_driver.rs b/crates/dbx-core/src/db/influxdb_driver.rs index 234ae865f..d73e30168 100644 --- a/crates/dbx-core/src/db/influxdb_driver.rs +++ b/crates/dbx-core/src/db/influxdb_driver.rs @@ -347,7 +347,7 @@ pub async fn list_databases(client: &InfluxdbClient) -> Result return Ok(influx_v2_buckets(client, Duration::from_secs(30)) .await? .into_iter() - .map(|bucket| DatabaseInfo { name: bucket.name }) + .map(|bucket| DatabaseInfo { name: bucket.name, size_bytes: None }) .collect()); } let result = influx_query(client, "SHOW DATABASES", None).await?; @@ -356,7 +356,7 @@ pub async fn list_databases(client: &InfluxdbClient) -> Result .iter() .flat_map(|r| &r.series) .flat_map(|s| &s.values) - .map(|row| DatabaseInfo { name: row[0].as_str().unwrap_or("").to_string() }) + .map(|row| DatabaseInfo { name: row[0].as_str().unwrap_or("").to_string(), size_bytes: None }) .collect()) } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 7252319a6..83aecb0e0 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -17,7 +17,7 @@ use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, CompletionAssistantResponse, DatabaseInfo, - ForeignKeyInfo, IndexInfo, ObjectInfo, ObjectStatistics, QueryResult, TableInfo, TriggerInfo, + DatabaseStatistics, ForeignKeyInfo, IndexInfo, ObjectInfo, ObjectStatistics, QueryResult, TableInfo, TriggerInfo, }; use super::file_validator::validate_file_path; @@ -1417,9 +1417,12 @@ pub async fn connect_bare_with_pool_limit_and_setup_database( } pub async fn list_databases(pool: &MySqlPool) -> Result, String> { + list_databases_without_size(pool).await +} + +async fn list_databases_without_size(pool: &MySqlPool) -> Result, String> { let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; - let result = match conn.query_iter("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME").await - { + let result = match conn.query_iter(list_databases_sql()).await { Ok(result) => result, Err(err) => { log::debug!("Falling back to SHOW DATABASES after information_schema.SCHEMATA failed: {err}"); @@ -1444,6 +1447,280 @@ pub async fn list_databases_show(pool: &MySqlPool) -> Result, Ok(database_infos_from_names(rows.iter().map(|row| get_str(row, 0)), true)) } +fn list_databases_sql() -> &'static str { + "SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME" +} + +/// Upper bound (ms) for server-side execution of MySQL statistics queries so a +/// slow or busy server cannot hang the sidebar. Emitted as an optimizer hint, +/// which MySQL (5.7.8+) honors and MariaDB / other forks silently ignore. Kept +/// below the frontend statistics timeout so the server aborts a stuck query +/// before the client gives up, instead of leaving it running server-side. +const MYSQL_STATISTICS_MAX_EXECUTION_TIME_MS: u32 = 8_000; + +/// Client-side execution bound for the information_schema size/statistics scans. +/// The `MAX_EXECUTION_TIME` hint above is ignored by MariaDB and other forks, so +/// without this a slow scan there would hold the pooled connection until the +/// outer metadata timeout fires — starving interactive queries (e.g. opening a +/// table) and making the UI feel stuck. This bound works on every server. +const MYSQL_STATISTICS_QUERY_TIMEOUT_SECS: u64 = 8; + +/// Loads a statistics result set on an existing connection, bounded by `timeout`. +/// Returns `Ok(None)` when the query does not finish in time, in which case the +/// connection must not be reused (its result stream is abandoned). +async fn mysql_query_statistics_rows_within( + conn: &mut mysql_async::Conn, + sql: &str, + timeout: Duration, +) -> Result>, String> { + if timeout.is_zero() { + return Ok(None); + } + match tokio::time::timeout(timeout, async { + let result = conn.query_iter(sql).await.map_err(|e| e.to_string())?; + result.collect_and_drop::().await.map_err(|e| e.to_string()) + }) + .await + { + Ok(Ok(rows)) => Ok(Some(rows)), + Ok(Err(err)) => Err(err), + Err(_) => Ok(None), + } +} + +/// Loads a bounded statistics result set using the default per-query budget. +async fn mysql_query_statistics_rows_bounded( + conn: &mut mysql_async::Conn, + sql: &str, +) -> Result>, String> { + mysql_query_statistics_rows_within(conn, sql, Duration::from_secs(MYSQL_STATISTICS_QUERY_TIMEOUT_SECS)).await +} + +pub async fn list_database_statistics(pool: &MySqlPool) -> Result, String> { + let databases = list_databases(pool).await?; + let database_names: Vec = databases.into_iter().map(|database| database.name).collect(); + if database_names.is_empty() { + return Ok(vec![]); + } + + let has_global_statistics_privilege = match mysql_current_user_has_global_database_statistics_privilege(pool).await + { + Ok(value) => value, + Err(err) => { + log::debug!("Unable to read MySQL global statistics privileges; database sizes will remain NULL: {err}"); + false + } + }; + let schema_statistics_privileges = if has_global_statistics_privilege { + HashSet::new() + } else { + match mysql_current_user_schema_statistics_privileges(pool).await { + Ok(value) => value, + Err(err) => { + log::debug!( + "Unable to read MySQL schema statistics privileges; database sizes will remain NULL: {err}" + ); + HashSet::new() + } + } + }; + + let databases_with_full_statistics: Vec = database_names + .iter() + .filter(|name| has_global_statistics_privilege || schema_statistics_privileges.contains(*name)) + .cloned() + .collect(); + let size_by_database = if databases_with_full_statistics.is_empty() { + HashMap::new() + } else { + match mysql_database_size_by_schema(pool, &databases_with_full_statistics).await { + Ok(value) => value, + Err(err) => { + log::debug!("Unable to read MySQL database statistics; database sizes will remain NULL: {err}"); + HashMap::new() + } + } + }; + + Ok(database_names + .into_iter() + .map(|name| { + let has_full_schema_statistics = + has_global_statistics_privilege || schema_statistics_privileges.contains(&name); + DatabaseStatistics { + size_bytes: mysql_database_statistics_size( + has_full_schema_statistics, + size_by_database.get(&name).copied(), + ), + name, + } + }) + .collect()) +} + +async fn mysql_current_user_has_global_database_statistics_privilege(pool: &MySqlPool) -> Result { + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; + let result = conn.query_iter(mysql_global_database_statistics_privilege_sql()).await.map_err(|e| e.to_string())?; + let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; + Ok(rows.first().and_then(|row| get_opt_i64(row, "HAS_PRIVILEGE")).unwrap_or(0) != 0) +} + +pub(crate) fn mysql_global_database_statistics_privilege_sql() -> &'static str { + "SELECT CASE WHEN EXISTS ( \ + SELECT 1 \ + FROM information_schema.USER_PRIVILEGES \ + WHERE REPLACE(GRANTEE, '''', '') = CURRENT_USER() \ + AND PRIVILEGE_TYPE IN ('SELECT','INSERT','UPDATE','DELETE','REFERENCES','INDEX','ALTER') \ + ) THEN 1 ELSE 0 END AS HAS_PRIVILEGE" +} + +async fn mysql_current_user_schema_statistics_privileges(pool: &MySqlPool) -> Result, String> { + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; + let result = conn.query_iter(mysql_schema_database_statistics_privileges_sql()).await.map_err(|e| e.to_string())?; + let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; + Ok(rows + .iter() + .map(|row| get_str_by_name(row, "TABLE_SCHEMA").trim().to_string()) + .filter(|schema| !schema.is_empty()) + .collect()) +} + +pub(crate) fn mysql_schema_database_statistics_privileges_sql() -> &'static str { + "SELECT TABLE_SCHEMA \ + FROM information_schema.SCHEMA_PRIVILEGES \ + WHERE REPLACE(GRANTEE, '''', '') = CURRENT_USER() \ + AND PRIVILEGE_TYPE IN ('SELECT','INSERT','UPDATE','DELETE','REFERENCES','INDEX','ALTER') \ + GROUP BY TABLE_SCHEMA" +} + +async fn mysql_database_size_by_schema( + pool: &MySqlPool, + database_names: &[String], +) -> Result, String> { + // Small batches let each completed query commit its schemas' sizes before the + // shared time budget runs out. On servers where scanning every schema at once + // exceeds the budget (returning nothing), this yields partial sizes for the + // schemas that did respond instead of all-or-nothing. + const BATCH_SIZE: usize = 8; + let mut sizes = HashMap::new(); + // Reuse a single connection for the whole statistics pass: the session-level + // stats setting below only applies to the connection it runs on, and reusing + // one connection avoids repeated checkout latency across batches. + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; + // Prefer cached information_schema statistics (MySQL 8.0+) so DATA_LENGTH / + // INDEX_LENGTH come from the data dictionary instead of being recomputed from + // the storage engine on every call. Recomputation is the main cause of slow + // or hanging size queries on servers with many tables. Best-effort: forks + // without this variable (MariaDB, MySQL 5.7) reject it but stay usable, which + // is expected — log at trace level so it is not mistaken for a real error. + if let Err(err) = conn.query_drop(mysql_cached_statistics_session_sql()).await { + log::trace!("MySQL information_schema_stats_expiry unsupported (using live stats): {err}"); + } + // Wall-clock budget shared across all batches so the whole pass (not each + // batch) stays bounded and releases the pooled connection promptly. + let budget = Duration::from_secs(MYSQL_STATISTICS_QUERY_TIMEOUT_SECS); + let started_at = Instant::now(); + for batch in database_names.chunks(BATCH_SIZE) { + if batch.is_empty() { + continue; + } + let remaining = budget.checked_sub(started_at.elapsed()).unwrap_or_default(); + let sql = mysql_database_size_by_schema_sql(batch); + let rows = match mysql_query_statistics_rows_within(&mut conn, &sql, remaining).await? { + Some(rows) => rows, + None => { + log::debug!( + "MySQL database size scan hit the {}s budget after {} schemas; remaining sizes left unknown", + MYSQL_STATISTICS_QUERY_TIMEOUT_SECS, + sizes.len() + ); + break; + } + }; + for row in rows { + let name = get_str_by_name(&row, "TABLE_SCHEMA").trim().to_string(); + if name.is_empty() { + continue; + } + if let Some(size_bytes) = get_opt_i64(&row, "SIZE_BYTES") { + sizes.insert(name, size_bytes); + } + } + } + Ok(sizes) +} + +fn mysql_cached_statistics_session_sql() -> &'static str { + "SET SESSION information_schema_stats_expiry = 86400" +} + +pub(crate) fn mysql_database_size_by_schema_sql(database_names: &[String]) -> String { + let database_filter = database_names.iter().map(|name| quote_value(name)).collect::>().join(", "); + format!( + "SELECT /*+ MAX_EXECUTION_TIME({max_ms}) */ TABLE_SCHEMA, \ + SUM(COALESCE(DATA_LENGTH, 0) + COALESCE(INDEX_LENGTH, 0)) AS SIZE_BYTES \ + FROM information_schema.TABLES \ + WHERE TABLE_TYPE <> 'VIEW' \ + AND TABLE_SCHEMA IN ({database_filter}) \ + GROUP BY TABLE_SCHEMA", + max_ms = MYSQL_STATISTICS_MAX_EXECUTION_TIME_MS, + ) +} + +pub(crate) fn mysql_database_statistics_size( + has_full_schema_statistics: bool, + visible_size_bytes: Option, +) -> Option { + if has_full_schema_statistics { + visible_size_bytes + } else { + None + } +} + +/// Computes the size of a single database on demand (manual right-click action). +/// +/// Scanning `information_schema.TABLES` for one schema is far cheaper than the +/// all-database automatic pass, so callers can afford a longer `timeout`. Keeps +/// the same NULL-safety as the automatic path: without full statistics privilege +/// on the schema (or when the bounded query does not finish) it returns `None` +/// rather than a misleading partial or zero size. +pub async fn database_size(pool: &MySqlPool, database: &str, timeout: Duration) -> Result, String> { + let has_full_schema_statistics = match mysql_current_user_has_global_database_statistics_privilege(pool).await { + Ok(true) => true, + Ok(false) => mysql_current_user_schema_statistics_privileges(pool) + .await + .map(|schemas| schemas.contains(database)) + .unwrap_or(false), + Err(err) => { + log::debug!("Unable to read MySQL statistics privileges; database size will remain NULL: {err}"); + false + } + }; + if !has_full_schema_statistics { + return Ok(None); + } + + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; + if let Err(err) = conn.query_drop(mysql_cached_statistics_session_sql()).await { + log::trace!("MySQL information_schema_stats_expiry unsupported (using live stats): {err}"); + } + let sql = mysql_database_size_by_schema_sql(&[database.to_string()]); + let rows = match mysql_query_statistics_rows_within(&mut conn, &sql, timeout).await? { + Some(rows) => rows, + None => return Ok(None), + }; + let size_bytes = rows.iter().find_map(|row| { + let name = get_str_by_name(row, "TABLE_SCHEMA").trim().to_string(); + if name == database { + get_opt_i64(row, "SIZE_BYTES") + } else { + None + } + }); + Ok(size_bytes) +} + fn database_infos_from_names( names: impl IntoIterator, include_catalogless_when_blank: bool, @@ -1454,12 +1731,12 @@ fn database_infos_from_names( .filter_map(|name| { saw_row = true; let name = name.trim().to_string(); - (!name.is_empty()).then_some(DatabaseInfo { name }) + (!name.is_empty()).then_some(DatabaseInfo { name, size_bytes: None }) }) .collect(); databases.sort_by(|a, b| a.name.cmp(&b.name)); if databases.is_empty() && saw_row && include_catalogless_when_blank { - return vec![DatabaseInfo { name: String::new() }]; + return vec![DatabaseInfo { name: String::new(), size_bytes: None }]; } databases } @@ -2140,17 +2417,32 @@ pub async fn list_objects( Ok(PagedObjectList { objects, paging_applied }) } -pub async fn list_object_statistics(pool: &MySqlPool, database: &str) -> Result, String> { - let sql = format!( - "SELECT TABLE_NAME, TABLE_ROWS, COALESCE(DATA_LENGTH, 0) + COALESCE(INDEX_LENGTH, 0) AS TOTAL_BYTES \ +fn mysql_object_statistics_sql(database: &str) -> String { + format!( + "SELECT /*+ MAX_EXECUTION_TIME({max_ms}) */ TABLE_NAME, TABLE_ROWS, \ + COALESCE(DATA_LENGTH, 0) + COALESCE(INDEX_LENGTH, 0) AS TOTAL_BYTES \ FROM information_schema.TABLES \ - WHERE TABLE_SCHEMA = {} AND TABLE_TYPE <> 'VIEW' \ + WHERE TABLE_SCHEMA = {schema} AND TABLE_TYPE <> 'VIEW' \ ORDER BY TABLE_NAME", - quote_value(database), - ); + max_ms = MYSQL_STATISTICS_MAX_EXECUTION_TIME_MS, + schema = quote_value(database), + ) +} + +pub async fn list_object_statistics(pool: &MySqlPool, database: &str) -> Result, String> { + let sql = mysql_object_statistics_sql(database); let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; - let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; - let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; + if let Err(err) = conn.query_drop(mysql_cached_statistics_session_sql()).await { + log::trace!("MySQL information_schema_stats_expiry unsupported (using live stats): {err}"); + } + let rows = match mysql_query_statistics_rows_bounded(&mut conn, &sql).await? { + Some(rows) => rows, + None => { + return Err(format!( + "MySQL object statistics query for `{database}` exceeded {MYSQL_STATISTICS_QUERY_TIMEOUT_SECS}s" + )) + } + }; Ok(rows .iter() .filter_map(|row| { @@ -4033,6 +4325,75 @@ mod tests { assert!(no_marker.is_empty()); } + #[test] + fn mysql_list_databases_sql_does_not_load_size_statistics() { + let sql = list_databases_sql(); + + assert!(sql.contains("information_schema.SCHEMATA")); + assert!(!sql.contains("information_schema.TABLES")); + assert!(!sql.contains("DATA_LENGTH")); + assert!(!sql.contains("INDEX_LENGTH")); + } + + #[test] + fn mysql_database_statistics_sql_preserves_missing_sizes_and_ignores_table_only_privileges() { + let global_privilege_sql = mysql_global_database_statistics_privilege_sql(); + let schema_privilege_sql = mysql_schema_database_statistics_privileges_sql(); + let size_sql = mysql_database_size_by_schema_sql(&["app".to_string(), "we'ird".to_string()]); + + assert!(global_privilege_sql.contains("information_schema.USER_PRIVILEGES")); + assert!(schema_privilege_sql.contains("information_schema.SCHEMA_PRIVILEGES")); + assert!(schema_privilege_sql.contains("GROUP BY TABLE_SCHEMA")); + assert!(size_sql.contains("information_schema.TABLES")); + assert!(size_sql.contains("DATA_LENGTH")); + assert!(size_sql.contains("INDEX_LENGTH")); + assert!(size_sql.contains("SIZE_BYTES")); + assert!(size_sql.contains("TABLE_TYPE <> 'VIEW'")); + assert!(size_sql.contains("TABLE_SCHEMA IN ('app', 'we\\'ird')")); + assert!(!global_privilege_sql.contains("information_schema.TABLE_PRIVILEGES")); + assert!(!schema_privilege_sql.contains("information_schema.TABLE_PRIVILEGES")); + assert!(!size_sql.contains("information_schema.TABLE_PRIVILEGES")); + assert!(!size_sql.contains("COALESCE(SIZE_BYTES, 0)")); + } + + #[test] + fn mysql_database_size_query_is_bounded_and_uses_cached_statistics() { + let size_sql = mysql_database_size_by_schema_sql(&["app".to_string()]); + + // Server-side execution bound so a slow information_schema scan cannot hang the sidebar. + assert!(size_sql.contains(&format!("MAX_EXECUTION_TIME({MYSQL_STATISTICS_MAX_EXECUTION_TIME_MS})"))); + // Bound stays under the 10s frontend statistics timeout. + assert!(MYSQL_STATISTICS_MAX_EXECUTION_TIME_MS < 10_000); + // Prefer cached data-dictionary statistics over live storage-engine recomputation. + assert_eq!(mysql_cached_statistics_session_sql(), "SET SESSION information_schema_stats_expiry = 86400"); + // The manual single-database size action reuses this builder with one schema. + assert!(size_sql.contains("TABLE_SCHEMA IN ('app')")); + } + + #[test] + fn mysql_object_statistics_query_is_bounded() { + let sql = mysql_object_statistics_sql("app"); + + // Object-browser size/row scan is the same expensive information_schema.TABLES query, + // so it must carry the same server-side execution bound to avoid stalling table opening. + assert!(sql.contains(&format!("MAX_EXECUTION_TIME({MYSQL_STATISTICS_MAX_EXECUTION_TIME_MS})"))); + assert!(sql.contains("information_schema.TABLES")); + assert!(sql.contains("TABLE_SCHEMA = 'app'")); + assert!(sql.contains("TABLE_TYPE <> 'VIEW'")); + // Client-side bound is universal (covers MariaDB / MySQL 5.7 that ignore the hint) and + // stays under the 12s backend statistics timeout in schema.rs. + assert!(MYSQL_STATISTICS_QUERY_TIMEOUT_SECS < 12); + } + + #[test] + fn mysql_database_statistics_size_keeps_null_for_low_privilege_or_missing_stats() { + assert_eq!(mysql_database_statistics_size(false, Some(4096)), None); + assert_eq!(mysql_database_statistics_size(false, None), None); + assert_eq!(mysql_database_statistics_size(true, None), None); + assert_eq!(mysql_database_statistics_size(true, Some(0)), Some(0)); + assert_eq!(mysql_database_statistics_size(true, Some(4096)), Some(4096)); + } + #[test] fn mysql_show_metadata_sql_supports_catalogless_services() { assert_eq!(show_tables_sql("", true), "SHOW FULL TABLES"); diff --git a/crates/dbx-core/src/db/ob_oracle.rs b/crates/dbx-core/src/db/ob_oracle.rs index cd1b5d226..3bc6b5862 100644 --- a/crates/dbx-core/src/db/ob_oracle.rs +++ b/crates/dbx-core/src/db/ob_oracle.rs @@ -40,7 +40,7 @@ pub async fn list_databases(pool: &mysql_async::Pool) -> Result = result.collect_and_drop().await.map_err(|e| e.to_string())?; - Ok(rows.iter().map(|row| DatabaseInfo { name: get_str(row, 0) }).collect()) + Ok(rows.iter().map(|row| DatabaseInfo { name: get_str(row, 0), size_bytes: None }).collect()) } pub async fn list_schemas(pool: &mysql_async::Pool) -> Result, String> { diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index a95603ab2..25e9a4d45 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -26,8 +26,8 @@ use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, CompletionAssistantResponse, DatabaseInfo, - ExtensionInfo, ForeignKeyInfo, FunctionInfo, IndexInfo, ObjectInfo, ObjectStatistics, OwnerInfo, QueryResult, - RuleInfo, SchemaInfo, SequenceInfo, TableInfo, TriggerInfo, + DatabaseStatistics, ExtensionInfo, ForeignKeyInfo, FunctionInfo, IndexInfo, ObjectInfo, ObjectStatistics, + OwnerInfo, QueryResult, RuleInfo, SchemaInfo, SequenceInfo, TableInfo, TriggerInfo, }; fn pg_temporal_to_json_value(row: &Row, idx: usize) -> Option { @@ -1401,17 +1401,60 @@ fn validate_postgres_ssl_paths(url: &str) -> Result<(), String> { pub async fn list_databases(pool: &Pool) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; - let rows = postgres_query_cached( - &client, - "SELECT datname FROM pg_database \ - WHERE datallowconn = true \ - ORDER BY datname", - &[], - ) - .await - .map_err(|e| e.to_string())?; + let rows = postgres_query_cached(&client, list_databases_sql(), &[]).await.map_err(|e| e.to_string())?; + + Ok(rows.iter().map(|row| DatabaseInfo { name: pg_row_try_string(row, 0), size_bytes: None }).collect()) +} + +fn list_databases_sql() -> &'static str { + "SELECT datname \ + FROM pg_database \ + WHERE datallowconn = true \ + ORDER BY datname" +} + +pub async fn list_database_statistics(pool: &Pool) -> Result, String> { + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let rows = postgres_query_cached(&client, list_database_statistics_sql(), &[]).await.map_err(|e| e.to_string())?; + + Ok(rows + .iter() + .map(|row| DatabaseStatistics { + name: pg_row_try_string(row, 0), + size_bytes: row.try_get::<_, Option>(1).ok().flatten(), + }) + .collect()) +} - Ok(rows.iter().map(|row| DatabaseInfo { name: pg_row_try_string(row, 0) }).collect()) +fn list_database_statistics_sql() -> &'static str { + "SELECT datname, \ + CASE WHEN has_database_privilege(datname, 'CONNECT') \ + THEN pg_database_size(oid) \ + ELSE NULL END AS size_bytes \ + FROM pg_database \ + WHERE datallowconn = true \ + ORDER BY datname" +} + +/// Computes the size of a single database on demand (manual right-click action). +/// +/// Mirrors [`list_database_statistics`] but scoped to one `datname`, so the +/// caller can apply a longer timeout without scanning every database. Keeps the +/// `has_database_privilege(..., 'CONNECT')` guard so a low-privilege account +/// gets `None` instead of failing the query. +pub async fn database_size(pool: &Pool, database: &str) -> Result, String> { + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let rows = postgres_query_cached(&client, database_size_sql(), &[&database]).await.map_err(|e| e.to_string())?; + Ok(rows.first().and_then(|row| row.try_get::<_, Option>(0).ok().flatten())) +} + +fn database_size_sql() -> &'static str { + "SELECT CASE WHEN has_database_privilege(datname, 'CONNECT') \ + THEN pg_database_size(oid) \ + ELSE NULL END AS size_bytes \ + FROM pg_database \ + WHERE datallowconn = true \ + AND datname = $1" } pub async fn list_tables(pool: &Pool, schema: &str) -> Result, String> { @@ -3282,6 +3325,31 @@ mod tests { use std::time::Instant; use tokio_postgres::types::FromSql; + #[test] + fn postgres_list_databases_sql_does_not_load_database_sizes() { + let sql = list_databases_sql(); + + assert!(sql.contains("pg_database")); + assert!(!sql.contains("pg_database_size")); + } + + #[test] + fn postgres_database_statistics_sql_loads_sizes_separately() { + let sql = list_database_statistics_sql(); + + assert!(sql.contains("pg_database_size")); + assert!(sql.contains("has_database_privilege")); + } + + #[test] + fn postgres_database_size_sql_scopes_to_one_database_with_privilege_guard() { + let sql = database_size_sql(); + + assert!(sql.contains("pg_database_size")); + assert!(sql.contains("has_database_privilege")); + assert!(sql.contains("datname = $1")); + } + struct DockerPostgres { name: String, port: u16, diff --git a/crates/dbx-core/src/db/rqlite_driver.rs b/crates/dbx-core/src/db/rqlite_driver.rs index f7ace1e45..388a00b83 100644 --- a/crates/dbx-core/src/db/rqlite_driver.rs +++ b/crates/dbx-core/src/db/rqlite_driver.rs @@ -98,7 +98,7 @@ pub async fn test_connection(client: &RqliteClient, timeout: Duration) -> Result } pub async fn list_databases(_client: &RqliteClient) -> Result, String> { - Ok(vec![DatabaseInfo { name: "main".to_string() }]) + Ok(vec![DatabaseInfo { name: "main".to_string(), size_bytes: None }]) } pub async fn list_tables(client: &RqliteClient, _schema: &str) -> Result, String> { diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index 34ef2dc3d..46f61df7f 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -1081,7 +1081,7 @@ mod tests { } pub async fn list_databases(_pool: &SqliteHandle) -> Result, String> { - Ok(vec![DatabaseInfo { name: "main".to_string() }]) + Ok(vec![DatabaseInfo { name: "main".to_string(), size_bytes: None }]) } pub async fn list_tables(pool: &SqliteHandle, _schema: &str) -> Result, String> { diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index d5157e5d3..9734ca7ec 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -760,7 +760,10 @@ pub async fn list_databases(client: &mut SqlServerClient) -> Result(0).unwrap_or("").to_string() }).collect()) + Ok(rows + .iter() + .map(|row| DatabaseInfo { name: row.get::<&str, _>(0).unwrap_or("").to_string(), size_bytes: None }) + .collect()) } pub async fn test_connection(client: &mut SqlServerClient) -> Result<(), String> { @@ -805,7 +808,7 @@ pub async fn list_linked_server_catalogs( Ok(rows .iter() .filter_map(|row| row.get::<&str, _>(0).map(str::trim).filter(|name| !name.is_empty())) - .map(|name| DatabaseInfo { name: name.to_string() }) + .map(|name| DatabaseInfo { name: name.to_string(), size_bytes: None }) .collect()) } diff --git a/crates/dbx-core/src/db/turso_driver.rs b/crates/dbx-core/src/db/turso_driver.rs index aff76be5b..1c06b9e64 100644 --- a/crates/dbx-core/src/db/turso_driver.rs +++ b/crates/dbx-core/src/db/turso_driver.rs @@ -137,7 +137,7 @@ pub async fn test_connection(client: &TursoClient, timeout: Duration) -> Result< } pub async fn list_databases(_client: &TursoClient) -> Result, String> { - Ok(vec![DatabaseInfo { name: "main".to_string() }]) + Ok(vec![DatabaseInfo { name: "main".to_string(), size_bytes: None }]) } pub async fn list_tables(client: &TursoClient, _schema: &str) -> Result, String> { diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 5e19fe760..d51be57e7 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -2,11 +2,18 @@ use crate::connection::{connection_url_for_endpoint, database_connection_config, use crate::db; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query::{agent_execute_query_params, should_discard_pool_after_error, QueryExecutionOptions}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::sync::Arc; use std::time::{Duration, Instant}; +const DATABASE_STATISTICS_TIMEOUT: Duration = Duration::from_secs(12); + +/// Timeout for the manual, single-database size request (right-click action). +/// Scanning one schema is much cheaper than the automatic all-database pass, so +/// this can be far more generous — slow servers still finish for one database. +const MANUAL_DATABASE_SIZE_TIMEOUT: Duration = Duration::from_secs(60); + macro_rules! extract_pool { ($connections:expr, $key:expr, $variant:ident) => { $connections.get($key).and_then(|v| match v { @@ -109,7 +116,7 @@ pub fn duckdb_list_databases_with_attached( let rows = stmt .query_map([], |row| { let name = row.get::<_, String>(0)?; - Ok(db::DatabaseInfo { name: if name == primary { "main".to_string() } else { name } }) + Ok(db::DatabaseInfo { name: if name == primary { "main".to_string() } else { name }, size_bytes: None }) }) .map_err(|e| e.to_string())?; Ok(rows.filter_map(|row| row.ok()).collect()) @@ -627,6 +634,120 @@ pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Resul retry_metadata_connection(state, connection_id, None, || list_databases_once(state, connection_id)).await } +pub async fn list_database_statistics_core( + state: &AppState, + connection_id: &str, +) -> Result, String> { + let start = Instant::now(); + let result = tokio::time::timeout( + DATABASE_STATISTICS_TIMEOUT, + retry_metadata_connection(state, connection_id, None, || list_database_statistics_once(state, connection_id)), + ) + .await; + match result { + Ok(Ok(statistics)) => { + log::info!( + "[list_database_statistics:done] connection_id={} count={} elapsed_ms={}", + connection_id, + statistics.len(), + start.elapsed().as_millis() + ); + Ok(statistics) + } + Ok(Err(error)) => { + log::warn!( + "[list_database_statistics:failed] connection_id={} elapsed_ms={} error={}", + connection_id, + start.elapsed().as_millis(), + error + ); + Err(error) + } + Err(_) => { + log::warn!( + "[list_database_statistics:timeout] connection_id={} timeout_ms={}", + connection_id, + DATABASE_STATISTICS_TIMEOUT.as_millis() + ); + Ok(Vec::new()) + } + } +} + +/// Computes the size of a single database on demand (manual right-click action). +/// +/// Unlike [`list_database_statistics_core`], which scans every database under a +/// short budget and leaves sizes blank on timeout, this scopes to one database +/// and allows [`MANUAL_DATABASE_SIZE_TIMEOUT`]. On timeout it returns `Ok(None)` +/// so the caller can surface "unavailable" without treating it as an error. +pub async fn database_size_core(state: &AppState, connection_id: &str, database: &str) -> Result, String> { + enum SizePool { + Mysql(db::mysql::MySqlPool), + Postgres(deadpool_postgres::Pool), + Unsupported, + } + + let db_config = connection_config(state, connection_id).await; + + // JDBC/Agent-backed MySQL uses a separate code path (no native pool). Apply + // the same permission/completeness protection as the native driver. + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Mysql) { + let connections = state.connections.read().await; + if let Some(client) = extract_pool!(&connections, connection_id, Agent) { + drop(connections); + let current_database = db_config.as_ref().and_then(|config| config.database.as_deref()); + let agent_timeout = agent_metadata_timeout(db_config.as_ref()); + return match tokio::time::timeout( + MANUAL_DATABASE_SIZE_TIMEOUT, + mysql_agent_database_size(client, database, current_database, agent_timeout), + ) + .await + { + Ok(inner) => inner, + Err(_) => { + log::warn!( + "[database_size:timeout] connection_id={connection_id} timeout_ms={} (agent mysql)", + MANUAL_DATABASE_SIZE_TIMEOUT.as_millis() + ); + Ok(None) + } + }; + } + } + + let pool = { + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Connection not found")? { + PoolKind::Mysql(_, mode) if *mode == MysqlMode::OceanBaseOracle => SizePool::Unsupported, + PoolKind::Mysql(_, _) if db_config.as_ref().is_some_and(is_doris_family_config) => SizePool::Unsupported, + PoolKind::Mysql(pool, _) => SizePool::Mysql(pool.clone()), + PoolKind::Postgres(pool) => SizePool::Postgres(pool.clone()), + _ => SizePool::Unsupported, + } + }; + + let database = database.to_string(); + let result = tokio::time::timeout(MANUAL_DATABASE_SIZE_TIMEOUT, async move { + match pool { + SizePool::Mysql(pool) => db::mysql::database_size(&pool, &database, MANUAL_DATABASE_SIZE_TIMEOUT).await, + SizePool::Postgres(pool) => db::postgres::database_size(&pool, &database).await, + SizePool::Unsupported => Ok(None), + } + }) + .await; + + match result { + Ok(inner) => inner, + Err(_) => { + log::warn!( + "[database_size:timeout] connection_id={connection_id} timeout_ms={}", + MANUAL_DATABASE_SIZE_TIMEOUT.as_millis() + ); + Ok(None) + } + } +} + pub async fn list_sqlserver_linked_servers_core( state: &AppState, connection_id: &str, @@ -886,7 +1007,7 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result Result Result Result, String> { + log::info!("[list_database_statistics] connection_id={connection_id}"); + let db_config = connection_config(state, connection_id).await; + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Mysql) { + let connections = state.connections.read().await; + if let Some(client) = extract_pool!(&connections, connection_id, Agent) { + drop(connections); + return mysql_agent_list_database_statistics( + client, + db_config.as_ref().and_then(|config| config.database.as_deref()), + agent_metadata_timeout(db_config.as_ref()), + ) + .await; + } + } + enum StatisticsPool { + Mysql(db::mysql::MySqlPool), + Postgres(deadpool_postgres::Pool), + Unsupported, + } + + // Clone the pool handle before awaiting database I/O. Holding the + // connections read lock here blocks creation of tab-scoped query pools, + // which can make opening a table spin until this optional query finishes. + let pool = { + let connections = state.connections.read().await; + match connections.get(connection_id).ok_or("Connection not found")? { + PoolKind::Mysql(_, mode) if *mode == MysqlMode::OceanBaseOracle => StatisticsPool::Unsupported, + PoolKind::Mysql(_, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { + StatisticsPool::Unsupported + } + PoolKind::Mysql(pool, _) => StatisticsPool::Mysql(pool.clone()), + PoolKind::Postgres(pool) => StatisticsPool::Postgres(pool.clone()), + _ => StatisticsPool::Unsupported, + } + }; + + match pool { + StatisticsPool::Mysql(pool) => db::mysql::list_database_statistics(&pool).await, + StatisticsPool::Postgres(pool) => db::postgres::list_database_statistics(&pool).await, + StatisticsPool::Unsupported => Ok(vec![]), + } +} + +async fn mysql_agent_execute_statistics_rows( + client: &mut db::agent_driver::AgentDriverClient, + sql: &str, + database: Option<&str>, + timeout_duration: Option, +) -> Result { + let params = agent_execute_query_params( + sql, + database, + None, + QueryExecutionOptions { max_rows: Some(10_000), timeout_secs: Some(10), ..Default::default() }, + ); + client.execute_query_with_timeout::(params, timeout_duration).await +} + +/// Detects whether the agent's current MySQL user has full statistics visibility. +/// +/// Mirrors the native driver's protection: `information_schema.TABLES` only +/// returns rows for objects the user can see, so without this a user with +/// partial table privileges would have a truncated sum reported as the full +/// database size. Returns `(has_global, per_schema_privileged)`. Best-effort: +/// on query failure we assume no privilege, which keeps sizes NULL rather than +/// reporting an incomplete value. +async fn mysql_agent_statistics_privileges( + client: &mut db::agent_driver::AgentDriverClient, + database: Option<&str>, + timeout_duration: Option, +) -> (bool, HashSet) { + let has_global = match mysql_agent_execute_statistics_rows( + client, + db::mysql::mysql_global_database_statistics_privilege_sql(), + database, + timeout_duration, + ) + .await + { + Ok(result) => result.rows.first().and_then(|row| row.first()).and_then(json_value_as_i64).unwrap_or(0) != 0, + Err(err) => { + log::debug!("Unable to read MySQL agent global statistics privileges; sizes will remain NULL: {err}"); + false + } + }; + if has_global { + return (true, HashSet::new()); + } + let schema_privileges = match mysql_agent_execute_statistics_rows( + client, + db::mysql::mysql_schema_database_statistics_privileges_sql(), + database, + timeout_duration, + ) + .await + { + Ok(result) => result + .rows + .into_iter() + .filter_map(|row| row.first().and_then(json_value_as_string).map(|schema| schema.trim().to_string())) + .filter(|schema| !schema.is_empty()) + .collect(), + Err(err) => { + log::debug!("Unable to read MySQL agent schema statistics privileges; sizes will remain NULL: {err}"); + HashSet::new() + } + }; + (false, schema_privileges) +} + +async fn mysql_agent_list_database_statistics( + client: Arc>, + database: Option<&str>, + timeout_duration: Option, +) -> Result, String> { + let mut client = client.lock().await; + let (has_global, schema_privileges) = + mysql_agent_statistics_privileges(&mut client, database, timeout_duration).await; + let result = mysql_agent_execute_statistics_rows( + &mut client, + mysql_agent_database_statistics_sql(), + database, + timeout_duration, + ) + .await?; + let raw = mysql_agent_database_statistics_from_query_result(result); + Ok(apply_mysql_agent_statistics_protection(raw, has_global, &schema_privileges)) +} + +/// Nulls out sizes for schemas the current user lacks full statistics privilege +/// on, so a privilege-truncated `information_schema.TABLES` sum is never +/// reported as a complete database size. Pure so it can be unit-tested. +fn apply_mysql_agent_statistics_protection( + raw: Vec, + has_global: bool, + schema_privileges: &HashSet, +) -> Vec { + raw.into_iter() + .map(|statistics| { + let has_full = has_global || schema_privileges.contains(&statistics.name); + db::DatabaseStatistics { + size_bytes: db::mysql::mysql_database_statistics_size(has_full, statistics.size_bytes), + name: statistics.name, + } + }) + .collect() +} + +/// Single-database size via the agent (manual right-click), with the same +/// completeness protection as [`mysql_agent_list_database_statistics`]. +async fn mysql_agent_database_size( + client: Arc>, + database: &str, + current_database: Option<&str>, + timeout_duration: Option, +) -> Result, String> { + let mut client = client.lock().await; + let (has_global, schema_privileges) = + mysql_agent_statistics_privileges(&mut client, current_database, timeout_duration).await; + if !(has_global || schema_privileges.contains(database)) { + return Ok(None); + } + let sql = db::mysql::mysql_database_size_by_schema_sql(&[database.to_string()]); + let result = mysql_agent_execute_statistics_rows(&mut client, &sql, current_database, timeout_duration).await?; + Ok(mysql_agent_database_statistics_from_query_result(result) + .into_iter() + .find(|statistics| statistics.name == database) + .and_then(|statistics| statistics.size_bytes)) +} + +fn mysql_agent_database_statistics_sql() -> &'static str { + "SELECT TABLE_SCHEMA AS name, \ + SUM(COALESCE(DATA_LENGTH, 0) + COALESCE(INDEX_LENGTH, 0)) AS size_bytes \ + FROM information_schema.TABLES \ + WHERE TABLE_TYPE <> 'VIEW' \ + GROUP BY TABLE_SCHEMA \ + ORDER BY TABLE_SCHEMA" +} + +fn mysql_agent_database_statistics_from_query_result(result: db::QueryResult) -> Vec { + result + .rows + .into_iter() + .filter_map(|row| { + let name = row.first().and_then(json_value_as_string)?.trim().to_string(); + if name.is_empty() { + return None; + } + let size_bytes = row.get(1).and_then(json_value_as_i64); + Some(db::DatabaseStatistics { name, size_bytes }) + }) + .collect() +} + +fn json_value_as_string(value: &serde_json::Value) -> Option<&str> { + match value { + serde_json::Value::String(value) => Some(value), + _ => None, + } +} + +fn json_value_as_i64(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::Number(value) => { + value.as_i64().or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok())) + } + serde_json::Value::String(value) => value.trim().parse::().ok(), + _ => None, + } +} + pub async fn list_schemas_core(state: &AppState, connection_id: &str, database: &str) -> Result, String> { list_schemas_core_with_visible_filter(state, connection_id, database, false).await } @@ -2415,16 +2751,17 @@ fn escape_presto_like_pattern(value: &str) -> String { mod tests { use super::db; use super::{ - clickhouse_metadata_database, deduplicate_column_infos, filter_mysql_system_databases_for_config, - filter_object_infos, filter_table_infos, filter_visible_schema_names, - is_agent_postgres_metadata_fallback_config, is_retryable_metadata_error, mysql_object_source_sql, - mysql_table_metadata_catalog, normalize_information_schema_table_type, oracle_columns_from_query_result, - oracle_columns_sql, oracle_object_statistics_dba_segments_sql, oracle_object_statistics_from_query_result, - oracle_object_statistics_rows_only_sql, oracle_object_statistics_sql, - oracle_object_statistics_user_segments_sql, oracle_table_comment_from_query_result, oracle_table_comment_sql, - oracle_table_comments_from_query_result, oracle_table_comments_sql, presto_like_columns_from_query_result, - presto_like_information_schema_columns_sql, presto_like_information_schema_tables_sql, - presto_like_tables_from_query_result, visible_schema_filter, + apply_mysql_agent_statistics_protection, clickhouse_metadata_database, deduplicate_column_infos, + filter_mysql_system_databases_for_config, filter_object_infos, filter_table_infos, filter_visible_schema_names, + is_agent_postgres_metadata_fallback_config, is_retryable_metadata_error, + mysql_agent_database_statistics_from_query_result, mysql_agent_database_statistics_sql, + mysql_object_source_sql, mysql_table_metadata_catalog, normalize_information_schema_table_type, + oracle_columns_from_query_result, oracle_columns_sql, oracle_object_statistics_dba_segments_sql, + oracle_object_statistics_from_query_result, oracle_object_statistics_rows_only_sql, + oracle_object_statistics_sql, oracle_object_statistics_user_segments_sql, + oracle_table_comment_from_query_result, oracle_table_comment_sql, oracle_table_comments_from_query_result, + oracle_table_comments_sql, presto_like_columns_from_query_result, presto_like_information_schema_columns_sql, + presto_like_information_schema_tables_sql, presto_like_tables_from_query_result, visible_schema_filter, }; #[cfg(feature = "duckdb-bundled")] use super::{ @@ -2432,7 +2769,7 @@ mod tests { duckdb_query_tables_in_database, }; use crate::models::connection::{ConnectionConfig, DatabaseType}; - use std::collections::HashMap; + use std::collections::{HashMap, HashSet}; fn test_column(name: &str, comment: Option<&str>, is_primary_key: bool) -> super::db::ColumnInfo { super::db::ColumnInfo { @@ -2510,6 +2847,66 @@ mod tests { assert_eq!(mysql_table_metadata_catalog("app_db", "tenant_db"), "tenant_db"); } + #[test] + fn mysql_agent_database_statistics_maps_query_result() { + let sql = mysql_agent_database_statistics_sql(); + assert!(sql.contains("information_schema.TABLES")); + assert!(sql.contains("DATA_LENGTH")); + assert!(sql.contains("INDEX_LENGTH")); + + let result = db::QueryResult { + columns: vec!["name".to_string(), "size_bytes".to_string()], + column_types: vec![], + column_sortables: vec![], + rows: vec![ + vec![serde_json::json!("app_db"), serde_json::json!(4096)], + vec![serde_json::json!("metrics"), serde_json::json!("8192")], + vec![serde_json::json!(""), serde_json::json!(1)], + vec![serde_json::json!("overflow"), serde_json::json!("not-a-number")], + ], + affected_rows: 0, + execution_time_ms: 0, + truncated: false, + session_id: None, + has_more: false, + }; + + let statistics = mysql_agent_database_statistics_from_query_result(result); + assert_eq!(statistics.len(), 3); + assert_eq!(statistics[0].name, "app_db"); + assert_eq!(statistics[0].size_bytes, Some(4096)); + assert_eq!(statistics[1].name, "metrics"); + assert_eq!(statistics[1].size_bytes, Some(8192)); + assert_eq!(statistics[2].name, "overflow"); + assert_eq!(statistics[2].size_bytes, None); + } + + #[test] + fn mysql_agent_statistics_protection_nulls_sizes_without_full_privilege() { + let raw = vec![ + db::DatabaseStatistics { name: "app_db".to_string(), size_bytes: Some(4096) }, + db::DatabaseStatistics { name: "metrics".to_string(), size_bytes: Some(8192) }, + ]; + + // Global statistics privilege: every visible size is trusted as complete. + let global = apply_mysql_agent_statistics_protection(raw.clone(), true, &HashSet::new()); + assert_eq!(global[0].size_bytes, Some(4096)); + assert_eq!(global[1].size_bytes, Some(8192)); + + // Only per-schema privilege on "app_db": "metrics" may be a truncated + // partial-visibility sum, so it must be reported as unknown, not wrong. + let mut privileged = HashSet::new(); + privileged.insert("app_db".to_string()); + let scoped = apply_mysql_agent_statistics_protection(raw.clone(), false, &privileged); + assert_eq!(scoped[0].size_bytes, Some(4096)); + assert_eq!(scoped[1].size_bytes, None); + + // No statistics privilege at all: nothing is trusted. + let none = apply_mysql_agent_statistics_protection(raw, false, &HashSet::new()); + assert_eq!(none[0].size_bytes, None); + assert_eq!(none[1].size_bytes, None); + } + #[test] fn mysql_object_source_sql_qualifies_cross_database_objects() { assert_eq!( @@ -2610,7 +3007,7 @@ mod tests { } fn test_database_info(name: &str) -> super::db::DatabaseInfo { - super::db::DatabaseInfo { name: name.to_string() } + super::db::DatabaseInfo { name: name.to_string(), size_bytes: None } } #[test] diff --git a/crates/dbx-core/src/schema/duckdb_metadata.rs b/crates/dbx-core/src/schema/duckdb_metadata.rs index 9ef0a133b..c5765c0e0 100644 --- a/crates/dbx-core/src/schema/duckdb_metadata.rs +++ b/crates/dbx-core/src/schema/duckdb_metadata.rs @@ -62,7 +62,7 @@ pub fn duckdb_list_databases_with_attached( let rows = stmt .query_map([], |row| { let name = row.get::<_, String>(0)?; - Ok(db::DatabaseInfo { name: if name == primary { "main".to_string() } else { name } }) + Ok(db::DatabaseInfo { name: if name == primary { "main".to_string() } else { name }, size_bytes: None }) }) .map_err(|e| e.to_string())?; Ok(rows.filter_map(|row| row.ok()).collect()) diff --git a/crates/dbx-core/src/types.rs b/crates/dbx-core/src/types.rs index 108277443..df6713859 100644 --- a/crates/dbx-core/src/types.rs +++ b/crates/dbx-core/src/types.rs @@ -3,6 +3,14 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DatabaseInfo { pub name: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub size_bytes: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DatabaseStatistics { + pub name: String, + pub size_bytes: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 52a634af1..03db6ec51 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -2,6 +2,7 @@ use std::collections::BTreeSet; use std::sync::atomic::{AtomicBool, Ordering}; use dbx_core::connection::AppState; +use dbx_core::db::mysql::{execute_query_with_max_rows, list_database_statistics, list_databases}; use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; use dbx_core::query::{execute_multi_core, execute_sql_statement}; use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest}; @@ -78,6 +79,21 @@ fn json_cell_text(value: &serde_json::Value) -> String { } } +fn mysql_url_with_credentials(admin_url: &str, database: &str, user: &str, password: &str) -> String { + let opts = mysql_async::Opts::from_url(admin_url).expect("DBX_LIVE_MYSQL_ADMIN_URL should be a valid URL"); + let host = opts.ip_or_hostname(); + let port = opts.tcp_port(); + format!("mysql://{user}:{password}@{host}:{port}/{database}") +} + +fn mysql_quote_identifier(value: &str) -> String { + format!("`{}`", value.replace('`', "``")) +} + +fn mysql_quote_string(value: &str) -> String { + format!("'{}'", value.replace('\\', "\\\\").replace('\'', "''")) +} + #[tokio::test] #[ignore = "requires the remote DBX MySQL 5.7 smoke-test container"] async fn live_mysql57_text_protocol_select_succeeds() { @@ -98,6 +114,63 @@ async fn live_mysql57_text_protocol_select_succeeds() { assert_eq!(result.rows, vec![vec![serde_json::json!("1"), serde_json::json!("mysql57")]]); } +#[tokio::test] +#[ignore = "requires DBX_LIVE_MYSQL_ADMIN_URL for a writable MySQL endpoint with CREATE USER/GRANT privileges"] +async fn live_mysql_database_statistics_keeps_null_for_table_only_privileges() { + let admin_url = std::env::var("DBX_LIVE_MYSQL_ADMIN_URL").expect("DBX_LIVE_MYSQL_ADMIN_URL"); + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let database = format!("dbx_stats_{suffix}"); + let user = format!("dbx_stats_u_{}", &suffix[..16]); + let password = format!("dbxStats{}!", &suffix[..12]); + let quoted_database = mysql_quote_identifier(&database); + let quoted_user = mysql_quote_string(&user); + let quoted_password = mysql_quote_string(&password); + let admin_pool = dbx_core::db::mysql::connect(&admin_url, std::time::Duration::from_secs(10)).await.unwrap(); + + let setup_sql = [ + format!("DROP USER IF EXISTS {quoted_user}@'%'"), + format!("DROP DATABASE IF EXISTS {quoted_database}"), + format!("CREATE DATABASE {quoted_database}"), + format!("CREATE TABLE {quoted_database}.visible_table (id INT PRIMARY KEY, payload VARCHAR(32))"), + format!("CREATE TABLE {quoted_database}.hidden_table (id INT PRIMARY KEY, payload VARCHAR(32))"), + format!("INSERT INTO {quoted_database}.visible_table VALUES (1, 'visible')"), + format!("INSERT INTO {quoted_database}.hidden_table VALUES (1, 'hidden')"), + format!("CREATE USER {quoted_user}@'%' IDENTIFIED BY {quoted_password}"), + format!("GRANT SELECT ON {quoted_database}.visible_table TO {quoted_user}@'%'"), + ]; + for sql in setup_sql { + execute_query_with_max_rows(&admin_pool, &sql, false, Some(10), Default::default()) + .await + .unwrap_or_else(|err| panic!("setup SQL failed: {sql}: {err}")); + } + + let limited_url = mysql_url_with_credentials(&admin_url, &database, &user, &password); + let limited_pool = dbx_core::db::mysql::connect(&limited_url, std::time::Duration::from_secs(10)).await.unwrap(); + + let databases_result = list_databases(&limited_pool).await; + let statistics_result = list_database_statistics(&limited_pool).await; + + let cleanup_sql = + [format!("DROP USER IF EXISTS {quoted_user}@'%'"), format!("DROP DATABASE IF EXISTS {quoted_database}")]; + for sql in cleanup_sql { + let _ = execute_query_with_max_rows(&admin_pool, &sql, false, Some(10), Default::default()).await; + } + + let databases = databases_result.expect("low-privilege user should still list databases"); + assert!(databases.iter().any(|entry| entry.name == database)); + assert!( + databases.iter().all(|entry| entry.size_bytes.is_none()), + "base database list must not load size statistics" + ); + + let statistics = statistics_result.expect("table-only privileges must not make statistics loading fail"); + let target = statistics.iter().find(|entry| entry.name == database).expect("target database statistics entry"); + assert_eq!( + target.size_bytes, None, + "table-only privileges expose an incomplete table set, so DBX must not report a partial size" + ); +} + #[tokio::test] #[ignore = "requires a remote MySQL-compatible endpoint with a limited result-set query"] async fn live_mysql_compatible_limited_text_protocol_query_succeeds() { diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 502f00495..2e7dc5f05 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -273,6 +273,8 @@ async fn main() { .route("/agents/progress/{operationId}", get(routes::agents::agent_progress)) // Schema .route("/schema/databases", get(routes::schema::list_databases)) + .route("/schema/database-statistics", get(routes::schema::list_database_statistics)) + .route("/schema/database-size", get(routes::schema::database_size)) .route("/schema/doris/catalogs", get(routes::schema::list_doris_catalogs)) .route("/schema/doris/catalog-databases", get(routes::schema::list_doris_catalog_databases)) .route("/schema/sqlserver/linked-servers", get(routes::schema::list_sqlserver_linked_servers)) diff --git a/crates/dbx-web/src/routes/schema.rs b/crates/dbx-web/src/routes/schema.rs index 39c30a34d..9a0200da2 100644 --- a/crates/dbx-web/src/routes/schema.rs +++ b/crates/dbx-web/src/routes/schema.rs @@ -32,6 +32,25 @@ pub async fn list_databases( Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } +pub async fn list_database_statistics( + State(state): State>, + Query(q): Query, +) -> Result, AppError> { + let result = + dbx_core::schema::list_database_statistics_core(&state.app, &q.connection_id).await.map_err(AppError)?; + Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) +} + +pub async fn database_size( + State(state): State>, + Query(q): Query, +) -> Result, AppError> { + let database = q.database.as_deref().ok_or_else(|| AppError("database is required".to_string()))?; + let result = + dbx_core::schema::database_size_core(&state.app, &q.connection_id, database).await.map_err(AppError)?; + Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) +} + /// Resolve a non-internal catalog for dispatch to the Doris multi-catalog path. async fn external_doris_catalog(state: &Arc, connection_id: &str, catalog: Option<&str>) -> Option { dbx_core::schema::resolve_external_doris_catalog(&state.app, connection_id, catalog).await diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index 6b0fbc37d..8ec0bc9a1 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -19,6 +19,23 @@ pub async fn list_databases( dbx_core::schema::list_databases_core(&state, &connection_id).await } +#[tauri::command] +pub async fn list_database_statistics( + state: State<'_, Arc>, + connection_id: String, +) -> Result, String> { + dbx_core::schema::list_database_statistics_core(&state, &connection_id).await +} + +#[tauri::command] +pub async fn database_size( + state: State<'_, Arc>, + connection_id: String, + database: String, +) -> Result, String> { + dbx_core::schema::database_size_core(&state, &connection_id, &database).await +} + #[tauri::command] pub async fn list_doris_catalogs( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 542af0a21..781dce7a5 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -935,6 +935,8 @@ pub fn run() { commands::plugins::install_jdbc_plugin_local, commands::plugins::uninstall_jdbc_plugin, commands::schema::list_databases, + commands::schema::list_database_statistics, + commands::schema::database_size, commands::schema::list_doris_catalogs, commands::schema::list_doris_catalog_databases, commands::schema::list_sqlserver_linked_servers,