diff --git a/apps/desktop/src/lib/tableTree.ts b/apps/desktop/src/lib/tableTree.ts index 087fcbf8d..2d4aed2c1 100644 --- a/apps/desktop/src/lib/tableTree.ts +++ b/apps/desktop/src/lib/tableTree.ts @@ -78,6 +78,8 @@ function makeTableTreeEntry({ isExpanded: false, children: [], }; + if (normalizedParentSchema) node.partitionParentSchema = normalizedParentSchema; + if (normalizedParentName) node.partitionParentName = normalizedParentName; return { key: objectIdentityKey(objectType, schema, name), @@ -305,6 +307,98 @@ export function hasTablePartitionGroups(node: TreeNode): boolean { return partitionGroupChildren(node).length > 0; } +export function mergeTableTreePageChildren(currentChildren: TreeNode[], pageChildren: TreeNode[], connectionId: string, database: string): TreeNode[] { + const roots = [...currentChildren]; + const nodesByKey = new Map(); + const rootKeys = new Set(); + + const nodeKey = (node: TreeNode) => objectIdentityKey("TABLE", node.schema, node.label); + const collect = (nodes: readonly TreeNode[]) => { + for (const node of nodes) { + if (node.type === "table") { + nodesByKey.set(nodeKey(node), node); + } + collect(node.children ?? []); + collect(node.hiddenChildren?.filter((child) => !(node.children ?? []).includes(child)) ?? []); + } + }; + + collect(roots); + for (const node of roots) { + if (node.type === "table") rootKeys.add(nodeKey(node)); + } + + const ensurePartitionGroup = (parent: TreeNode): TreeNode => { + const existing = partitionGroupChildren(parent)[0]; + if (existing) return existing; + const group: TreeNode = { + id: `${parent.id}:__partitions`, + label: "tree.partitions", + type: "group-partitions", + connectionId, + database, + schema: parent.schema, + tableName: parent.label, + objectCount: 0, + isExpanded: false, + children: [], + }; + parent.children = [...(parent.children ?? []), group]; + parent.hiddenChildren = [...(parent.hiddenChildren ?? []), group]; + return group; + }; + + const addToParent = (parent: TreeNode, child: TreeNode) => { + const group = ensurePartitionGroup(parent); + const children = group.children ?? []; + if (!children.some((node) => nodeKey(node) === nodeKey(child))) { + group.children = sortDatabaseObjectsByName([...children, child], (node) => node.label); + group.objectCount = group.children.length; + } + }; + + const addNode = (node: TreeNode) => { + if (node.type !== "table") { + roots.push(node); + return; + } + + const key = nodeKey(node); + if (nodesByKey.has(key)) return; + + const parentName = node.partitionParentName; + const parentSchema = node.partitionParentSchema || node.schema; + const parentKey = parentName ? objectIdentityKey("TABLE", parentSchema, parentName) : ""; + const parent = parentKey ? nodesByKey.get(parentKey) : undefined; + nodesByKey.set(key, node); + if (parent && parent !== node) { + addToParent(parent, node); + return; + } + + if (!rootKeys.has(key)) { + roots.push(node); + rootKeys.add(key); + } + }; + + const flattenIncomingTables = (node: TreeNode): TreeNode[] => { + if (node.type !== "table") return [node]; + const descendants = partitionGroupChildren(node) + .flatMap((group) => group.children ?? []) + .flatMap(flattenIncomingTables); + node.children = (node.children ?? []).filter((child) => child.type !== "group-partitions"); + node.hiddenChildren = (node.hiddenChildren ?? []).filter((child) => child.type !== "group-partitions"); + return [node, ...descendants]; + }; + + for (const node of pageChildren.flatMap(flattenIncomingTables)) { + addNode(node); + } + + return sortDatabaseObjectsByName(roots, (node) => node.label); +} + export type DatabaseObjectTreeKind = SidebarObjectKind; function buildObjectTreeEntries({ nodeId, connectionId, database, schema, objects, objectType }: { nodeId: string; connectionId: string; database: string; schema?: string; objects: ObjectInfo[]; objectType: "TABLE" | "VIEW" | "MATERIALIZED_VIEW" }): TreeNode[] { diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index b91be2a71..2cde76187 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -36,6 +36,7 @@ import { buildTableTreeNodes, expandCachedObjectBrowserNodes, mergeTableInfosIntoObjects, + mergeTableTreePageChildren, objectGroupRefreshParentId, objectTypesForGroupNode, tablePartitionGroups, @@ -1570,7 +1571,8 @@ export const useConnectionStore = defineStore("connection", () => { pageSize: node.loadMore.pageSize, }); const currentChildren = withoutLoadMoreNodes(parent.children); - const nextChildren = page.hasMore ? [...currentChildren, ...page.children, buildLoadMoreNode(parent, page.nextOffset, node.loadMore.pageSize)] : [...currentChildren, ...page.children]; + const mergedChildren = mergeTableTreePageChildren(currentChildren, page.children, parent.connectionId, parent.database); + const nextChildren = page.hasMore ? [...mergedChildren, buildLoadMoreNode(parent, page.nextOffset, node.loadMore.pageSize)] : mergedChildren; parent.objectCount = (parent.objectCount ?? currentChildren.length) + page.objectCount; setChildren(parent, nextChildren); await savePersistedTreeChildren(objectGroupCacheKey(parent), nextChildren); diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 9692fe722..72bc08892 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -473,6 +473,8 @@ export interface TreeNode { objectCount?: number; loadedKeyCount?: number; totalKeyCount?: number; + partitionParentSchema?: string; + partitionParentName?: string; hiddenChildren?: TreeNode[]; savedSqlId?: string; savedSqlFolderId?: string; diff --git a/packages/app-tests/tableTree.test.ts b/packages/app-tests/tableTree.test.ts index c07ff694e..a1b06a099 100644 --- a/packages/app-tests/tableTree.test.ts +++ b/packages/app-tests/tableTree.test.ts @@ -1,6 +1,6 @@ import { test } from "vitest"; import assert from "node:assert/strict"; -import { buildGroupedObjectTreeNodes, buildObjectGroupPlaceholderNodes, buildSimpleObjectTreeNodes, buildTableTreeNodes, mergeTableInfosIntoObjects } from "../../apps/desktop/src/lib/tableTree.ts"; +import { buildGroupedObjectTreeNodes, buildObjectGroupPlaceholderNodes, buildSimpleObjectTreeNodes, buildTableTreeNodes, mergeTableInfosIntoObjects, mergeTableTreePageChildren } from "../../apps/desktop/src/lib/tableTree.ts"; import type { ObjectInfo, TableInfo, TreeNode } from "../../apps/desktop/src/types/database.ts"; function table(name: string, parent?: string): TableInfo { @@ -75,6 +75,36 @@ test("buildTableTreeNodes keeps partitions visible when their parent is not load ); }); +test("mergeTableTreePageChildren attaches later page partitions to loaded parents", () => { + const firstPage = buildTableTreeNodes({ + nodeId: "conn:app:public", + connectionId: "conn", + database: "app", + schema: "public", + tables: [table("events"), table("events_region_0", "events")], + }); + const secondPage = buildTableTreeNodes({ + nodeId: "conn:app:public", + connectionId: "conn", + database: "app", + schema: "public", + tables: [table("events_region_0_2026_01", "events_region_0")], + }); + + const merged = mergeTableTreePageChildren(firstPage, secondPage, "conn", "app"); + assert.deepEqual( + merged.map((node) => node.label), + ["events"], + ); + + const regionPartition = partitionGroup(merged[0])?.children?.[0]; + assert.equal(regionPartition?.label, "events_region_0"); + assert.deepEqual( + partitionGroup(regionPartition!)?.children?.map((node) => node.label), + ["events_region_0_2026_01"], + ); +}); + test("buildTableTreeNodes keeps sidebar tables in natural name order", () => { const nodes = buildTableTreeNodes({ nodeId: "conn:app:public",