Skip to content

Commit ea00c85

Browse files
Light2Darkclaude
andauthored
feat(datasources): only auto-expand search matches in the data sources tree (#9955)
## Summary Follow-up to #9845 (which superseded the closed #9874). No API or wire-protocol change — purely a frontend UX improvement on the existing `Database.schemas` model. Previously, typing any text into the data sources search box ran `setIsExpanded(hasSearch)` on **every** database and schema, expanding the entire tree open regardless of whether a branch had any matching tables. Now a database/schema row auto-expands during search **only when its already-loaded subtree contains a table whose name matches the query**. Crucially, deferred (not-yet-fetched) tables and child schemas are treated as non-matching, so searching never triggers a lazy catalog fetch — important for large/remote connections (Iceberg, etc.) where unconditional expansion would cause a fetch storm. Expansion is re-evaluated on each keystroke. This is a reimplementation, against the merged `Database.schemas` model, of the matched-subtree search behavior from the closed #9874. ## Changes - New helpers in `datasources/utils.ts`: - `schemaSubtreeMatchesSearch(schema, query)` — recurses over *resolved* tables/child schemas only. - `shouldExpandDatabaseForSearch(database, query)` — false when the schema list is deferred. - `datasources.tsx`: `DatabaseItem` and `SchemaNode` now drive auto-expansion via these helpers, tracking `prevSearchValue` so expansion follows the query. Removes the now-unused `hasSearch` plumbing through the tree. https://github.com/user-attachments/assets/c58cb882-aa88-4b74-9508-ebf8bb6c1dee ## Testing - New unit tests for both helpers, including the deferred-bucket cases (search must not match unfetched data). - `pnpm test src/components/datasources/` — 55 tests pass. ## Pre-Review Checklist - [x] Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it. - [x] Video or media evidence is provided for any visual changes (optional). ## Merge Checklist - [x] I have read the contributor guidelines. - [x] Tests have been added for the changes made. <!-- This is an auto-generated description by cubic. --> <a href="https://cubic.dev/pr/marimo-team/marimo/pull/9955?utm_source=github" target="_blank" rel="noopener noreferrer" data-no-image-dialog="true"><picture><source media="(prefers-color-scheme: dark)" srcset="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://www.cubic.dev/buttons/review-in-cubic-light.svg"><img alt="Review in cubic" src="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"></picture></a> <!-- End of auto-generated description by cubic. --> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent c9b403d commit ea00c85

4 files changed

Lines changed: 901 additions & 19 deletions

File tree

frontend/src/components/datasources/__tests__/utils.test.ts

Lines changed: 127 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,18 @@
33
import { describe, expect, it } from "vitest";
44
import type { SQLTableContext } from "@/core/datasets/data-source-connections";
55
import { DUCKDB_ENGINE } from "@/core/datasets/engines";
6-
import type { DataTable, DataTableColumn } from "@/core/kernel/messages";
7-
import { sqlCode, tableUniqueId } from "../utils";
6+
import type {
7+
Database,
8+
DatabaseSchema,
9+
DataTable,
10+
DataTableColumn,
11+
} from "@/core/kernel/messages";
12+
import {
13+
schemaSubtreeMatchesSearch,
14+
shouldExpandDatabaseForSearch,
15+
sqlCode,
16+
tableUniqueId,
17+
} from "../utils";
818

919
describe("sqlCode", () => {
1020
const mockTable: DataTable = {
@@ -519,3 +529,118 @@ describe("tableUniqueId", () => {
519529
);
520530
});
521531
});
532+
533+
function makeTable(name: string): DataTable {
534+
return {
535+
name,
536+
columns: [],
537+
source: "memory",
538+
source_type: "local",
539+
type: "table",
540+
engine: null,
541+
indexes: null,
542+
num_columns: null,
543+
num_rows: null,
544+
variable_name: null,
545+
primary_keys: null,
546+
};
547+
}
548+
549+
function makeSchema(opts: {
550+
name: string;
551+
tables: DataTable[];
552+
tables_resolved?: boolean;
553+
child_schemas?: DatabaseSchema[];
554+
child_schemas_resolved?: boolean;
555+
}): DatabaseSchema {
556+
return {
557+
name: opts.name,
558+
tables: opts.tables,
559+
tables_resolved: opts.tables_resolved ?? true,
560+
child_schemas: opts.child_schemas ?? [],
561+
child_schemas_resolved: opts.child_schemas_resolved ?? true,
562+
};
563+
}
564+
565+
function makeDatabase(
566+
name: string,
567+
schemas: DatabaseSchema[],
568+
schemas_resolved = true,
569+
): Database {
570+
return { name, dialect: "duckdb", schemas, schemas_resolved, engine: null };
571+
}
572+
573+
describe("schemaSubtreeMatchesSearch", () => {
574+
it("returns false for an empty query", () => {
575+
const schema = makeSchema({ name: "main", tables: [makeTable("users")] });
576+
expect(schemaSubtreeMatchesSearch(schema, "")).toBe(false);
577+
expect(schemaSubtreeMatchesSearch(schema, " ")).toBe(false);
578+
expect(schemaSubtreeMatchesSearch(schema, undefined)).toBe(false);
579+
});
580+
581+
it("matches a table name case-insensitively", () => {
582+
const schema = makeSchema({ name: "main", tables: [makeTable("Users")] });
583+
expect(schemaSubtreeMatchesSearch(schema, "user")).toBe(true);
584+
expect(schemaSubtreeMatchesSearch(schema, "orders")).toBe(false);
585+
});
586+
587+
it("matches a table in a resolved child schema", () => {
588+
const schema = makeSchema({
589+
name: "parent",
590+
tables: [],
591+
child_schemas: [
592+
makeSchema({ name: "child", tables: [makeTable("orders")] }),
593+
],
594+
});
595+
expect(schemaSubtreeMatchesSearch(schema, "orders")).toBe(true);
596+
});
597+
598+
it("ignores deferred tables so search never triggers a fetch", () => {
599+
const schema = makeSchema({
600+
name: "main",
601+
tables: [makeTable("users")],
602+
tables_resolved: false,
603+
});
604+
expect(schemaSubtreeMatchesSearch(schema, "users")).toBe(false);
605+
});
606+
607+
it("ignores deferred child schemas", () => {
608+
const schema = makeSchema({
609+
name: "parent",
610+
tables: [],
611+
child_schemas: [
612+
makeSchema({ name: "child", tables: [makeTable("orders")] }),
613+
],
614+
child_schemas_resolved: false,
615+
});
616+
expect(schemaSubtreeMatchesSearch(schema, "orders")).toBe(false);
617+
});
618+
});
619+
620+
describe("shouldExpandDatabaseForSearch", () => {
621+
it("returns false for an empty query", () => {
622+
const db = makeDatabase("memory", [
623+
makeSchema({ name: "main", tables: [makeTable("users")] }),
624+
]);
625+
expect(shouldExpandDatabaseForSearch(db, "")).toBe(false);
626+
expect(shouldExpandDatabaseForSearch(db, undefined)).toBe(false);
627+
});
628+
629+
it("expands when a loaded schema contains a matching table", () => {
630+
const db = makeDatabase("memory", [
631+
makeSchema({ name: "main", tables: [makeTable("users")] }),
632+
makeSchema({ name: "other", tables: [makeTable("orders")] }),
633+
]);
634+
expect(shouldExpandDatabaseForSearch(db, "user")).toBe(true);
635+
expect(shouldExpandDatabaseForSearch(db, "nomatch")).toBe(false);
636+
});
637+
638+
it("does not expand when the schema list itself is deferred", () => {
639+
const db = makeDatabase(
640+
"memory",
641+
[makeSchema({ name: "main", tables: [makeTable("users")] })],
642+
/* schemas_resolved */ false,
643+
);
644+
expect(shouldExpandDatabaseForSearch(db, "users")).toBe(false);
645+
});
646+
});

frontend/src/components/datasources/datasources.tsx

Lines changed: 44 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ import {
8383
areSchemasResolved,
8484
areTablesResolved,
8585
isSchemaless,
86+
schemaSubtreeMatchesSearch,
87+
shouldExpandDatabaseForSearch,
8688
sqlCode,
8789
tableUniqueId,
8890
} from "./utils";
@@ -360,7 +362,6 @@ export const DataSources: React.FC = () => {
360362
key={database.name}
361363
connection={connection}
362364
database={database}
363-
hasSearch={hasSearch}
364365
searchValue={searchValue}
365366
/>
366367
))}
@@ -435,7 +436,6 @@ interface DataSourceTree {
435436
dialect: string;
436437
engineName: string;
437438
databaseName: string;
438-
hasSearch: boolean;
439439
searchValue?: string;
440440
}
441441

@@ -470,17 +470,15 @@ function buildSqlTableContext(
470470
const DatabaseTree: React.FC<{
471471
connection: DataSourceConnection;
472472
database: Database;
473-
hasSearch: boolean;
474473
searchValue?: string;
475-
}> = ({ connection, database, hasSearch, searchValue }) => {
474+
}> = ({ connection, database, searchValue }) => {
476475
const tree = React.useMemo<DataSourceTree>(
477476
() => ({
478477
engineName: connection.name,
479478
databaseName: database.name,
480479
dialect: connection.dialect,
481480
defaultSchema: connection.default_schema,
482481
defaultDatabase: connection.default_database,
483-
hasSearch,
484482
searchValue,
485483
}),
486484
[
@@ -489,7 +487,6 @@ const DatabaseTree: React.FC<{
489487
connection.default_schema,
490488
connection.default_database,
491489
database.name,
492-
hasSearch,
493490
searchValue,
494491
],
495492
);
@@ -498,7 +495,7 @@ const DatabaseTree: React.FC<{
498495
<DatabaseItem
499496
engineName={connection.name}
500497
database={database}
501-
hasSearch={hasSearch}
498+
searchValue={searchValue}
502499
>
503500
<DataSourceTreeContext.Provider value={tree}>
504501
<SchemaList
@@ -513,18 +510,30 @@ const DatabaseTree: React.FC<{
513510
};
514511

515512
const DatabaseItem: React.FC<{
516-
hasSearch: boolean;
513+
searchValue?: string;
517514
engineName: string;
518515
database: Database;
519516
children: React.ReactNode;
520-
}> = ({ hasSearch, engineName, database, children }) => {
521-
const [isExpanded, setIsExpanded] = React.useState(false);
517+
}> = ({ searchValue, engineName, database, children }) => {
518+
// Auto-expand during search only when a loaded schema under this database
519+
// contains a matching table.
520+
const expandForSearch = shouldExpandDatabaseForSearch(database, searchValue);
521+
const [isExpanded, setIsExpanded] = React.useState(expandForSearch);
522522
const [isSelected, setIsSelected] = React.useState(false);
523-
const [prevHasSearch, setPrevHasSearch] = React.useState(hasSearch);
524-
525-
if (prevHasSearch !== hasSearch) {
526-
setPrevHasSearch(hasSearch);
527-
setIsExpanded(hasSearch);
523+
const [prevSearchValue, setPrevSearchValue] = React.useState(searchValue);
524+
const [prevExpandForSearch, setPrevExpandForSearch] =
525+
React.useState(expandForSearch);
526+
527+
// Re-sync when the query changes or when newly resolved data turns this
528+
// branch into a match under the same query (deferred subtrees resolve
529+
// asynchronously, so the decision can flip without the query changing).
530+
if (
531+
prevSearchValue !== searchValue ||
532+
prevExpandForSearch !== expandForSearch
533+
) {
534+
setPrevSearchValue(searchValue);
535+
setPrevExpandForSearch(expandForSearch);
536+
setIsExpanded(expandForSearch);
528537
}
529538

530539
return (
@@ -657,12 +666,30 @@ interface SchemaNodeProps {
657666
const SchemaNode: React.FC<SchemaNodeProps> = (props) => {
658667
const { schema, schemaPath, depth } = props;
659668
const tree = useDataSourceTree();
660-
const { databaseName, hasSearch, searchValue } = tree;
661-
const [isExpanded, setIsExpanded] = React.useState(hasSearch);
669+
const { databaseName, searchValue } = tree;
670+
// Auto-expand during search only when this schema's loaded subtree contains a
671+
// matching table.
672+
const expandForSearch = schemaSubtreeMatchesSearch(schema, searchValue);
673+
const [isExpanded, setIsExpanded] = React.useState(expandForSearch);
662674
const [isSelected, setIsSelected] = React.useState(false);
675+
const [prevSearchValue, setPrevSearchValue] = React.useState(searchValue);
676+
const [prevExpandForSearch, setPrevExpandForSearch] =
677+
React.useState(expandForSearch);
663678
const uniqueValue = `${databaseName}:${schemaPath.join(".")}`;
664679
const childSchemas = schema.child_schemas ?? [];
665680

681+
// Re-sync when the query changes or when newly resolved data turns this
682+
// subtree into a match under the same query (deferred tables/child schemas
683+
// resolve asynchronously, so the decision can flip without the query changing).
684+
if (
685+
prevSearchValue !== searchValue ||
686+
prevExpandForSearch !== expandForSearch
687+
) {
688+
setPrevSearchValue(searchValue);
689+
setPrevExpandForSearch(expandForSearch);
690+
setIsExpanded(expandForSearch);
691+
}
692+
666693
return (
667694
<>
668695
<CommandItem

frontend/src/components/datasources/utils.ts

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,51 @@ export function areChildSchemasResolved(schema: DatabaseSchema): boolean {
5252
return schema.child_schemas_resolved !== false;
5353
}
5454

55+
/**
56+
* Whether a loaded schema subtree contains a table whose name matches
57+
* `searchValue`. Deferred (not-yet-fetched) tables and child schemas are
58+
* treated as non-matching, so searching never triggers a lazy catalog fetch on
59+
* large/remote connections.
60+
*/
61+
export function schemaSubtreeMatchesSearch(
62+
schema: DatabaseSchema,
63+
searchValue: string | undefined,
64+
): boolean {
65+
const query = searchValue?.trim().toLowerCase();
66+
if (!query) {
67+
return false;
68+
}
69+
if (
70+
areTablesResolved(schema) &&
71+
schema.tables.some((table) => table.name.toLowerCase().includes(query))
72+
) {
73+
return true;
74+
}
75+
if (areChildSchemasResolved(schema)) {
76+
return (schema.child_schemas ?? []).some((child) =>
77+
schemaSubtreeMatchesSearch(child, query),
78+
);
79+
}
80+
return false;
81+
}
82+
83+
/**
84+
* Auto-expand a database row during search only when one of its loaded schemas
85+
* contains a matching table. Returns false when the schema list itself is
86+
* deferred.
87+
*/
88+
export function shouldExpandDatabaseForSearch(
89+
database: Database,
90+
searchValue: string | undefined,
91+
): boolean {
92+
if (!searchValue?.trim() || !areSchemasResolved(database)) {
93+
return false;
94+
}
95+
return database.schemas.some((schema) =>
96+
schemaSubtreeMatchesSearch(schema, searchValue),
97+
);
98+
}
99+
55100
interface SqlCodeFormatter {
56101
/**
57102
* Format the table path based on dialect-specific rules

0 commit comments

Comments
 (0)