@marimo-team/frontend 0.16.2 → 0.16.3

package/src/core/codemirror/language/languages/markdown.ts

@@ -103,6 +103,9 @@ export class MarkdownLanguageAdapter
     code: string,
     metadata: MarkdownLanguageAdapterMetadata,
   ): [string, number] {
+    // NB. Must be kept consistent with marimo/_convert/utils.py
+    // ::markdown_to_marimo
+
     // Empty string
     if (code === "") {
       // Need at least a space, otherwise the output will be 6 quotes
@@ -152,7 +155,7 @@ export class MarkdownLanguageAdapter
     const tree = pythonLanguage.parser.parse(pythonCode);
 
     // This is the exact match of mo.md() signature
-    const enterOrder: Array<{ match: string | RegExp; stop?: boolean }> = [
+    const enterOrder: { match: string | RegExp; stop?: boolean }[] = [
       { match: "Script" },
       { match: "ExpressionStatement" },
       { match: "CallExpression" },

package/src/core/codemirror/language/languages/sql/{validation-errors.ts → banner-validation-errors.ts}

@@ -29,20 +29,26 @@ export function clearSqlValidationError(cellId: CellId) {
   store.set(sqlValidationErrorsAtom, newErrors);
 }
 
+export function clearAllSqlValidationErrors() {
+  store.set(sqlValidationErrorsAtom, new Map<CellId, SQLValidationError>());
+}
+
 export function setSqlValidationError({
   cellId,
-  error,
+  errorMessage,
   dialect,
 }: {
   cellId: CellId;
-  error: string;
+  errorMessage: string;
   dialect: SupportedDialects | null;
 }) {
   const sqlValidationErrors = store.get(sqlValidationErrorsAtom);
   const newErrors = new Map(sqlValidationErrors);
 
   const errorResult: SQLValidationError =
-    dialect === "DuckDB" ? handleDuckdbError(error) : splitErrorMessage(error);
+    dialect === "DuckDB"
+      ? handleDuckdbError(errorMessage)
+      : splitErrorMessage(errorMessage);
 
   newErrors.set(cellId, errorResult);
   store.set(sqlValidationErrorsAtom, newErrors);

package/src/core/codemirror/language/languages/sql/completion-builder.ts

@@ -0,0 +1,160 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Completion } from "@codemirror/autocomplete";
+import type { SQLNamespace } from "@codemirror/lang-sql";
+import { createRoot } from "react-dom/client";
+import type {
+  Database,
+  DatabaseSchema,
+  DataTable,
+  DataTableColumn,
+} from "@/core/kernel/messages";
+import {
+  renderColumnInfo,
+  renderDatabaseInfo,
+  renderSchemaInfo,
+  renderTableInfo,
+} from "./renderers";
+
+/**
+ * Simple builder for SQL completion schemas.
+ */
+export class CompletionBuilder {
+  private schema: Record<string, SQLNamespace> = {};
+
+  /**
+   * Add a table with its columns at the specified path
+   */
+  addTable(path: string[], table: DataTable): this {
+    const tableNamespace: SQLNamespace = {
+      self: tableToCompletion({
+        table: table,
+      }),
+      children: table.columns.map((col) =>
+        columnToCompletion({
+          column: col,
+        }),
+      ),
+    };
+
+    this.setAt([...path, table.name], tableNamespace);
+    return this;
+  }
+
+  /**
+   * Add a schema at the specified path
+   */
+  addSchema(path: string[], schema: DatabaseSchema): this {
+    const schemaObject: SQLNamespace = {
+      self: schemaToCompletion({
+        namespace: schema,
+        path: path,
+      }),
+      children: {},
+    };
+
+    this.setAt(path, schemaObject);
+    return this;
+  }
+
+  /**
+   * Add a database at the specified path
+   */
+  addDatabase(path: string[], database: Database): this {
+    const databaseObject: SQLNamespace = {
+      self: databaseToCompletion({
+        namespace: database,
+        path: path,
+      }),
+      children: {},
+    };
+    this.setAt(path, databaseObject);
+    return this;
+  }
+
+  /**
+   * Set a value at a nested path, creating intermediate objects as needed
+   */
+  private setAt(path: string[], value: SQLNamespace): void {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    let current: any = this.schema;
+    for (const key of path.slice(0, -1)) {
+      if (!current[key]) {
+        current[key] = { children: {} };
+      }
+      current = current[key].children;
+    }
+    current[path[path.length - 1]] = value;
+  }
+
+  /**
+   * Build the final schema
+   */
+  build(): SQLNamespace {
+    return this.schema;
+  }
+
+  /**
+   * Reset for reuse
+   */
+  reset(): this {
+    this.schema = {};
+    return this;
+  }
+}
+
+function columnToCompletion(opts: { column: DataTableColumn }): Completion {
+  return {
+    label: opts.column.name,
+    type: "column",
+    info: () => {
+      const dom = document.createElement("div");
+      createRoot(dom).render(renderColumnInfo(opts.column));
+      return { dom: dom };
+    },
+  };
+}
+
+function tableToCompletion(opts: { table: DataTable }): Completion {
+  return {
+    label: opts.table.name,
+    type: "table",
+    info: () => {
+      const dom = document.createElement("div");
+      createRoot(dom).render(renderTableInfo(opts.table));
+      return { dom: dom };
+    },
+  };
+}
+
+function schemaToCompletion(opts: {
+  namespace: DatabaseSchema;
+  path: string[];
+}): Completion {
+  return {
+    label: opts.namespace.name,
+    detail: opts.path.join("."),
+    type: "schema",
+    info: () => {
+      const dom = document.createElement("div");
+      createRoot(dom).render(renderSchemaInfo(opts.namespace));
+      return { dom: dom };
+    },
+  };
+}
+
+function databaseToCompletion(opts: {
+  namespace: Database;
+  path: string[];
+}): Completion {
+  return {
+    label: opts.namespace.name,
+    detail: opts.path.join("."),
+    type: "database",
+    info: () => {
+      const dom = document.createElement("div");
+      createRoot(dom).render(renderDatabaseInfo(opts.namespace));
+      return { dom: dom };
+    },
+  };
+}

package/src/core/codemirror/language/languages/sql/completion-sources.tsx

@@ -7,6 +7,7 @@ import {
 } from "@codemirror/lang-sql";
 import type { EditorState } from "@codemirror/state";
 import { DefaultSqlTooltipRenders } from "@marimo-team/codemirror-sql";
+import { once } from "@/utils/once";
 import { languageMetadataField } from "../../metadata";
 import { SCHEMA_CACHE } from "./completion-store";
 import type { SQLLanguageAdapterMetadata } from "./sql";
@@ -28,7 +29,12 @@ export function tablesCompletionSource(): CompletionSource {
       return null;
     }
 
-    return schemaCompletionSource(config)(ctx);
+    const completions = schemaCompletionSource(config)(ctx);
+    if (!completions) {
+      return null;
+    }
+
+    return completions;
   };
 }
 
@@ -84,7 +90,7 @@ export function customKeywordCompletionSource(): CompletionSource {
 }
 
 // e.g. lazily load keyword docs
-const getKeywordDocs = async (): Promise<Record<string, unknown>> => {
+const getKeywordDocs = once(async (): Promise<Record<string, unknown>> => {
   const keywords = await import(
     "@marimo-team/codemirror-sql/data/common-keywords.json"
   );
@@ -96,4 +102,4 @@ const getKeywordDocs = async (): Promise<Record<string, unknown>> => {
     ...keywords.default.keywords,
     ...duckdbKeywords.default.keywords,
   };
-};
+});

package/src/core/codemirror/language/languages/sql/completion-store.ts

@@ -1,6 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import type { SQLConfig, SQLDialect } from "@codemirror/lang-sql";
+import { atom } from "jotai";
 import { isSchemaless } from "@/components/datasources/utils";
 import { dataConnectionsMapAtom } from "@/core/datasets/data-source-connections";
 import type { ConnectionName } from "@/core/datasets/engines";
@@ -8,14 +9,22 @@ import { datasetTablesAtom } from "@/core/datasets/state";
 import type { DataSourceConnection } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
 import { LRUCache } from "@/utils/lru";
+import { CompletionBuilder } from "./completion-builder";
 import { guessDialect, ModifiedStandardSQL } from "./utils";
 
-type TableToCols = Record<string, string[]>;
-type Schemas = Record<string, TableToCols>;
 type CachedSchema = Pick<SQLConfig, "schema" | "defaultSchema"> & {
   shouldAddLocalTables: boolean;
 };
 
+const datasetTableCompletionsAtom = atom((get) => {
+  const tables = get(datasetTablesAtom);
+  const builder = new CompletionBuilder();
+  for (const table of tables) {
+    builder.addTable([], table);
+  }
+  return builder.build();
+});
+
 class SQLCompletionStore {
   private cache: LRUCache<DataSourceConnection, CachedSchema>;
 
@@ -33,81 +42,74 @@ class SQLCompletionStore {
   }
 
   private getConnectionSchema(connection: DataSourceConnection): CachedSchema {
-    const schemaMap: Record<string, TableToCols> = {};
-    const databaseMap: Record<string, Schemas> = {};
+    const { default_database, databases, default_schema } = connection;
+    const builder = new CompletionBuilder();
 
     // When there is only one database, it is the default
-    const defaultDb = connection.databases.find(
-      (db) =>
-        db.name === connection.default_database ||
-        connection.databases.length === 1,
+    const defaultDb = databases.find(
+      (db) => db.name === default_database || databases.length === 1,
     );
 
-    const dbToVerify = defaultDb ?? connection.databases[0];
+    const dbToVerify = defaultDb ?? databases[0];
     const isSchemalessDb =
       dbToVerify?.schemas.some((schema) => isSchemaless(schema.name)) ?? false;
 
     // For schemaless databases, treat databases as schemas
     if (isSchemalessDb) {
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      const dbToTablesMap: Record<string, any> = {};
-
-      for (const db of connection.databases) {
+      for (const db of databases) {
         const isDefaultDb = db.name === defaultDb?.name;
-
-        for (const schema of db.schemas) {
-          for (const table of schema.tables) {
-            const columns = table.columns.map((col) => col.name);
-
-            if (isDefaultDb) {
-              // For default database, add tables directly to top level
-              dbToTablesMap[table.name] = columns;
-            } else {
-              // Otherwise nest under database name
-              dbToTablesMap[db.name] = dbToTablesMap[db.name] || {};
-              dbToTablesMap[db.name][table.name] = columns;
-            }
+        const tables = db.schemas.flatMap((schema) => schema.tables);
+        builder.addDatabase([db.name], db);
+
+        for (const table of tables) {
+          if (isDefaultDb) {
+            // For default database, add tables directly to top level
+            builder.addTable([], table);
+          } else {
+            // Otherwise nest under database name
+            builder.addTable([db.name], table);
           }
         }
       }
 
       return {
         shouldAddLocalTables: false,
-        schema: dbToTablesMap,
+        schema: builder.build(),
         defaultSchema: defaultDb?.name,
       };
     }
 
-    // For default db, we can use the schema name directly
-    for (const schema of defaultDb?.schemas ?? []) {
-      schemaMap[schema.name] = {};
-      for (const table of schema.tables) {
-        const columns = table.columns.map((col) => col.name);
-        schemaMap[schema.name][table.name] = columns;
+    // For default db, we can use the schema name directly so add them to the top level
+    if (defaultDb) {
+      for (const schema of defaultDb.schemas) {
+        builder.addSchema([schema.name], schema);
+
+        for (const table of schema.tables) {
+          builder.addTable([schema.name], table);
+        }
       }
     }
 
     // Otherwise, we need to use the fully qualified name
-    for (const database of connection.databases) {
-      if (database.name === defaultDb?.name) {
-        continue;
-      }
-      databaseMap[database.name] = {};
+    for (const database of databases) {
+      // We still want to add the default database here in case
+      // users want fully qualified names for completions
+
+      builder.addDatabase([database.name], database);
 
       for (const schema of database.schemas) {
-        databaseMap[database.name][schema.name] = {};
+        builder.addSchema([database.name, schema.name], schema);
 
         for (const table of schema.tables) {
-          const columns = table.columns.map((col) => col.name);
-          databaseMap[database.name][schema.name][table.name] = columns;
+          builder.addTable([database.name, schema.name], table);
         }
       }
     }
 
     return {
       shouldAddLocalTables: true,
-      schema: { ...databaseMap, ...schemaMap },
-      defaultSchema: connection.default_schema ?? undefined,
+      schema: builder.build(),
+      defaultSchema: default_schema ?? undefined,
     };
   }
 
@@ -142,15 +144,9 @@ class SQLCompletionStore {
     }
 
     const getTablesMap = () => {
-      const localTables = store.get(datasetTablesAtom);
       // If there is a conflict with connection tables,
       // the engine will prioritize the connection tables without special handling
-      const tablesMap: TableToCols = {};
-      for (const table of localTables) {
-        const tableColumns = table.columns.map((col) => col.name);
-        tablesMap[table.name] = tableColumns;
-      }
-      return tablesMap;
+      return store.get(datasetTableCompletionsAtom);
     };
 
     const schema = this.cache.getOrCreate(connection);