@marimo-team/frontend 0.14.18-dev29 → 0.14.18-dev30

package/src/core/codemirror/language/__tests__/sql.test.ts

@@ -6,6 +6,7 @@ import type {
 } from "@codemirror/autocomplete";
 import { PostgreSQL } from "@codemirror/lang-sql";
 import { EditorState, type Extension } from "@codemirror/state";
+import { DuckDBDialect } from "@marimo-team/codemirror-sql/dialects";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
 import type { DataSourceConnection } from "@/core/datasets/data-source-connections";
 import {
@@ -16,12 +17,11 @@ import { type ConnectionName, DUCKDB_ENGINE } from "@/core/datasets/engines";
 import { datasetsAtom } from "@/core/datasets/state";
 import type { DatasetsState } from "@/core/datasets/types";
 import { store } from "@/core/state/jotai";
+import { TestSQLCompletionStore } from "../languages/sql/completion-store";
 import {
-  SQLCompletionStore,
   SQLLanguageAdapter,
   type SQLLanguageAdapterMetadata,
-} from "../languages/sql";
-import { DuckDBDialect } from "../languages/sql-dialects/duckdb";
+} from "../languages/sql/sql";
 import { languageMetadataField } from "../metadata";
 
 const adapter = new SQLLanguageAdapter();
@@ -618,7 +618,7 @@ _df = mo.sql(
 
 describe("tablesCompletionSource", () => {
   const mockStore = store;
-  const completionStore = new SQLCompletionStore();
+  const completionStore = new TestSQLCompletionStore();
 
   beforeEach(() => {
     // Reset the adapter engine
@@ -1184,7 +1184,7 @@ describe("tablesCompletionSource", () => {
     ];
 
     describe("SQL Completions", () => {
-      const completionStore = new SQLCompletionStore();
+      const completionStore = new TestSQLCompletionStore();
 
       beforeEach(() => {
         // Reset state

package/src/core/codemirror/language/extension.ts

@@ -28,7 +28,7 @@ import { historyCompartment } from "../editing/extensions";
 import { formattingChangeEffect } from "../format";
 import { createPanel } from "../react-dom/createPanel";
 import { getLanguageAdapters, LanguageAdapters } from "./LanguageAdapters";
-import { initializeSQLDialect } from "./languages/sql";
+import { initializeSQLDialect } from "./languages/sql/sql";
 import type { LanguageMetadata } from "./metadata";
 import { languageMetadataField, setLanguageMetadata } from "./metadata";
 import { LanguagePanelComponent } from "./panel/panel";

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

@@ -0,0 +1,99 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Completion, CompletionSource } from "@codemirror/autocomplete";
+import {
+  keywordCompletionSource,
+  schemaCompletionSource,
+} from "@codemirror/lang-sql";
+import type { EditorState } from "@codemirror/state";
+import { DefaultSqlTooltipRenders } from "@marimo-team/codemirror-sql";
+import { languageMetadataField } from "../../metadata";
+import { SCHEMA_CACHE } from "./completion-store";
+import type { SQLLanguageAdapterMetadata } from "./sql";
+
+function getSQLMetadata(state: EditorState): SQLLanguageAdapterMetadata {
+  return state.field(languageMetadataField) as SQLLanguageAdapterMetadata;
+}
+
+/**
+ * Custom schema completion source that dynamically gets the Dialect and SQL tables.
+ */
+export function tablesCompletionSource(): CompletionSource {
+  return (ctx) => {
+    const metadata = getSQLMetadata(ctx.state);
+    const connectionName = metadata.engine;
+    const config = SCHEMA_CACHE.getCompletionSource(connectionName);
+
+    if (!config) {
+      return null;
+    }
+
+    return schemaCompletionSource(config)(ctx);
+  };
+}
+
+/**
+ * Custom keyword completion source that dynamically gets the Dialect.
+ * This also ignores keyword completions on table columns.
+ */
+export function customKeywordCompletionSource(): CompletionSource {
+  return (ctx) => {
+    const metadata = getSQLMetadata(ctx.state);
+    const connectionName = metadata.engine;
+    const dialect = SCHEMA_CACHE.getDialect(connectionName);
+
+    // We want to ignore keyword completions on something like
+    // `WHERE my_table.col`
+    //                    ^cursor
+    const textBefore = ctx.matchBefore(/\.\w*/);
+    if (textBefore) {
+      // If there is a match, we are typing after a dot,
+      // so we don't want to trigger SQL keyword completion
+      return null;
+    }
+
+    const keywordRenderer = (label: string, type: string): Completion => {
+      return {
+        label,
+        type,
+        info: async () => {
+          const keywordDocs = await getKeywordDocs();
+          const keywordInfo = keywordDocs[label.toLocaleLowerCase()];
+          if (!keywordInfo) {
+            return null;
+          }
+
+          const dom = document.createElement("div");
+          dom.innerHTML = DefaultSqlTooltipRenders.keyword({
+            keyword: label,
+            info: keywordInfo,
+          });
+          return dom;
+        },
+      };
+    };
+
+    const uppercaseKeywords = true;
+    const result = keywordCompletionSource(
+      dialect,
+      uppercaseKeywords,
+      keywordRenderer,
+    )(ctx);
+    return result;
+  };
+}
+
+// e.g. lazily load keyword docs
+const getKeywordDocs = async (): Promise<Record<string, unknown>> => {
+  const keywords = await import(
+    "@marimo-team/codemirror-sql/data/common-keywords.json"
+  );
+  // Include DuckDB for now, but we can remove this once we have a better way to handle dialect-specific keywords
+  const duckdbKeywords = await import(
+    "@marimo-team/codemirror-sql/data/duckdb-keywords.json"
+  );
+  return {
+    ...keywords.default.keywords,
+    ...duckdbKeywords.default.keywords,
+  };
+};

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

@@ -0,0 +1,191 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import {
+  type SQLConfig,
+  type SQLDialect,
+  StandardSQL,
+} from "@codemirror/lang-sql";
+import { isSchemaless } from "@/components/datasources/utils";
+import { dataConnectionsMapAtom } from "@/core/datasets/data-source-connections";
+import type { ConnectionName } from "@/core/datasets/engines";
+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 { guessDialect } from "./utils";
+
+type TableToCols = Record<string, string[]>;
+type Schemas = Record<string, TableToCols>;
+type CachedSchema = Pick<SQLConfig, "schema" | "defaultSchema"> & {
+  shouldAddLocalTables: boolean;
+};
+
+class SQLCompletionStore {
+  private cache: LRUCache<DataSourceConnection, CachedSchema>;
+
+  constructor() {
+    this.cache = new LRUCache(10, {
+      create: (connection) => this.getConnectionSchema(connection),
+    });
+  }
+
+  private getConnection(
+    connectionName: ConnectionName,
+  ): DataSourceConnection | undefined {
+    const dataConnectionsMap = store.get(dataConnectionsMapAtom);
+    return dataConnectionsMap.get(connectionName);
+  }
+
+  private getConnectionSchema(connection: DataSourceConnection): CachedSchema {
+    const schemaMap: Record<string, TableToCols> = {};
+    const databaseMap: Record<string, Schemas> = {};
+
+    // 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 dbToVerify = defaultDb ?? connection.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) {
+        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;
+            }
+          }
+        }
+      }
+
+      return {
+        shouldAddLocalTables: false,
+        schema: dbToTablesMap,
+        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;
+      }
+    }
+
+    // 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 schema of database.schemas) {
+        databaseMap[database.name][schema.name] = {};
+
+        for (const table of schema.tables) {
+          const columns = table.columns.map((col) => col.name);
+          databaseMap[database.name][schema.name][table.name] = columns;
+        }
+      }
+    }
+
+    return {
+      shouldAddLocalTables: true,
+      schema: { ...databaseMap, ...schemaMap },
+      defaultSchema: connection.default_schema ?? undefined,
+    };
+  }
+
+  /**
+   * Returns the raw dialect of the connection passed from the backend,
+   * or null if the connection is not found
+   */
+  getInternalDialect(connectionName: ConnectionName): string | null {
+    const connection = this.getConnection(connectionName);
+    if (!connection) {
+      return null;
+    }
+    return connection.dialect;
+  }
+
+  /**
+   * Get the inferred SQL dialect for a connection
+   * If the connection is not found, return the standard SQL dialect.
+   */
+  getDialect(connectionName: ConnectionName): SQLDialect {
+    const connection = this.getConnection(connectionName);
+    if (!connection) {
+      return StandardSQL;
+    }
+    return guessDialect(connection) ?? StandardSQL;
+  }
+
+  getCompletionSource(connectionName: ConnectionName): SQLConfig | null {
+    const connection = this.getConnection(connectionName);
+    if (!connection) {
+      return null;
+    }
+
+    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;
+    };
+
+    const schema = this.cache.getOrCreate(connection);
+
+    return {
+      dialect: guessDialect(connection),
+      schema: schema.shouldAddLocalTables
+        ? { ...schema.schema, ...getTablesMap() }
+        : schema.schema,
+      defaultSchema: schema.defaultSchema,
+      defaultTable: getSingleTable(connection),
+    };
+  }
+}
+
+function getSingleTable(connection: DataSourceConnection): string | undefined {
+  if (connection.databases.length !== 1) {
+    return undefined;
+  }
+  const database = connection.databases[0];
+  if (database.schemas.length !== 1) {
+    return undefined;
+  }
+  const schema = database.schemas[0];
+  if (schema.tables.length !== 1) {
+    return undefined;
+  }
+  return schema.tables[0].name;
+}
+
+export const SCHEMA_CACHE = new SQLCompletionStore();
+
+// For testing
+export { SQLCompletionStore as TestSQLCompletionStore };