generate different sql code based on dialect (#6915)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->

Fixes #6875. Might also tackle #6872.

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

Author: Light2Dark

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

@@ -214,6 +214,7 @@ export const DataSources: React.FC = () => {
                   databaseName={database.name}
                   hasSearch={hasSearch}
                   searchValue={searchValue}
+                  dialect={connection.dialect}
                 />
               </DatabaseItem>
             ))}
@@ -340,6 +341,7 @@ const SchemaList: React.FC<{
   schemas: DatabaseSchema[];
   defaultSchema?: string | null;
   defaultDatabase?: string | null;
+  dialect: string;
   engineName: string;
   databaseName: string;
   hasSearch: boolean;
@@ -348,6 +350,7 @@ const SchemaList: React.FC<{
   schemas,
   defaultSchema,
   defaultDatabase,
+  dialect,
   engineName,
   databaseName,
   hasSearch,
@@ -384,6 +387,7 @@ const SchemaList: React.FC<{
               schema: schema.name,
               defaultSchema: defaultSchema,
               defaultDatabase: defaultDatabase,
+              dialect: dialect,
             }}
           />
         </SchemaItem>

frontend/src/components/datasources/datasources.tsx

@@ -1,14 +1,90 @@
 /* Copyright 2024 Marimo. All rights reserved. */
+
+import { BigQueryDialect } from "@marimo-team/codemirror-sql/dialects";
+import { isKnownDialect } from "@/core/codemirror/language/languages/sql/utils";
 import type { SQLTableContext } from "@/core/datasets/data-source-connections";
 import { DUCKDB_ENGINE } from "@/core/datasets/engines";
 import type { DataTable, DataType } from "@/core/kernel/messages";
+import { logNever } from "@/utils/assertNever";
 import type { ColumnHeaderStatsKey } from "../data-table/types";
 
 // Some databases have no schemas, so we don't show it (eg. Clickhouse)
 export function isSchemaless(schemaName: string) {
   return schemaName === "";
 }
 
+interface SqlCodeFormatter {
+  /**
+   * Format the table name based on dialect-specific rules
+   */
+  formatTableName: (tableName: string) => string;
+  /**
+   * Format the SELECT clause
+   */
+  formatSelectClause: (columnName: string, tableName: string) => string;
+}
+
+const defaultFormatter: SqlCodeFormatter = {
+  formatTableName: (tableName: string) => tableName,
+  formatSelectClause: (columnName: string, tableName: string) =>
+    `SELECT ${columnName} FROM ${tableName} LIMIT 100`,
+};
+
+function getFormatter(dialect: string): SqlCodeFormatter {
+  dialect = dialect.toLowerCase();
+  if (!isKnownDialect(dialect)) {
+    return defaultFormatter;
+  }
+
+  switch (dialect) {
+    case "bigquery": {
+      const quote = BigQueryDialect.spec.identifierQuotes;
+      return {
+        // BigQuery uses backticks for identifiers
+        formatTableName: (tableName: string) => `${quote}${tableName}${quote}`,
+        formatSelectClause: defaultFormatter.formatSelectClause,
+      };
+    }
+    case "mssql":
+    case "sqlserver":
+      return {
+        formatTableName: defaultFormatter.formatTableName,
+        formatSelectClause: (columnName: string, tableName: string) =>
+          `SELECT TOP 100 ${columnName} FROM ${tableName}`,
+      };
+    case "timescaledb":
+      return {
+        // TimescaleDB uses double quotes for identifiers
+        formatTableName: (tableName: string) => {
+          const parts = tableName.split(".");
+          return parts.map((part) => `"${part}"`).join(".");
+        },
+        formatSelectClause: defaultFormatter.formatSelectClause,
+      };
+    case "postgresql":
+    case "postgres":
+    case "db2":
+    case "mysql":
+    case "sqlite":
+    case "duckdb":
+    case "mariadb":
+    case "cassandra":
+    case "noql":
+    case "athena":
+    case "hive":
+    case "redshift":
+    case "snowflake":
+    case "flink":
+    case "mongodb":
+    case "oracle":
+    case "oracledb":
+      return defaultFormatter;
+    default:
+      logNever(dialect);
+      return defaultFormatter;
+  }
+}
+
 export function sqlCode({
   table,
   columnName,
@@ -19,8 +95,14 @@ export function sqlCode({
   sqlTableContext?: SQLTableContext;
 }) {
   if (sqlTableContext) {
-    const { engine, schema, defaultSchema, defaultDatabase, database } =
-      sqlTableContext;
+    const {
+      engine,
+      schema,
+      defaultSchema,
+      defaultDatabase,
+      database,
+      dialect,
+    } = sqlTableContext;
     let tableName = table.name;
 
     // Set the fully qualified table name based on schema and database
@@ -39,11 +121,18 @@ export function sqlCode({
       }
     }
 
+    const formatter = getFormatter(dialect);
+    const formattedTableName = formatter.formatTableName(tableName);
+    const selectClause = formatter.formatSelectClause(
+      columnName,
+      formattedTableName,
+    );
+
     if (engine === DUCKDB_ENGINE) {
-      return `_df = mo.sql(f"SELECT ${columnName} FROM ${tableName} LIMIT 100")`;
+      return `_df = mo.sql(f"""\n${selectClause}\n""")`;
     }
 
-    return `_df = mo.sql(f"SELECT ${columnName} FROM ${tableName} LIMIT 100", engine=${engine})`;
+    return `_df = mo.sql(f"""\n${selectClause}\n""", engine=${engine})`;
   }
 
   return `_df = mo.sql(f'SELECT "${columnName}" FROM ${table.name} LIMIT 100')`;

frontend/src/components/datasources/utils.ts

@@ -134,7 +134,7 @@ class SQLCompletionStore {
     if (!connection) {
       return ModifiedStandardSQL;
     }
-    return guessDialect(connection) ?? ModifiedStandardSQL;
+    return guessDialect(connection);
   }
 
   getCompletionSource(connectionName: ConnectionName): SQLConfig | null {
@@ -152,7 +152,7 @@ class SQLCompletionStore {
     const schema = this.cache.getOrCreate(connection);
 
     return {
-      dialect: guessDialect(connection) ?? ModifiedStandardSQL,
+      dialect: guessDialect(connection),
       schema: schema.shouldAddLocalTables
         ? { ...schema.schema, ...getTablesMap() }
         : schema.schema,

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

@@ -38,6 +38,7 @@ import type { HotkeyProvider } from "@/core/hotkeys/hotkeys";
 import type { ValidateSQLResult } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
 import { resolvedThemeAtom } from "@/theme/useTheme";
+import { logNever } from "@/utils/assertNever";
 import { Logger } from "@/utils/Logger";
 import { variableCompletionSource } from "../../embedded/embedded-python";
 import { languageMetadataField } from "../../metadata";
@@ -52,6 +53,7 @@ import {
 } from "./completion-sources";
 import { SCHEMA_CACHE } from "./completion-store";
 import { getSQLMode, type SQLMode } from "./sql-mode";
+import { isKnownDialect } from "./utils";
 
 const DEFAULT_DIALECT = DuckDBDialect;
 const DEFAULT_PARSER_DIALECT = "DuckDB";
@@ -353,6 +355,11 @@ function connectionNameToParserDialect(
 ): ParserDialects | null {
   const dialect =
     SCHEMA_CACHE.getInternalDialect(connectionName)?.toLowerCase();
+
+  if (!dialect || !isKnownDialect(dialect)) {
+    return null;
+  }
+
   switch (dialect) {
     case "postgresql":
     case "postgres":
@@ -385,8 +392,15 @@ function connectionNameToParserDialect(
     case "flink":
       return "FlinkSQL";
     case "mongodb":
+    case "noql":
       return "Noql";
+    case "oracle":
+    case "oracledb":
+    case "timescaledb":
+      Logger.debug("Unsupported dialect", { dialect });
+      return null;
     default:
+      logNever(dialect);
       return null;
   }
 }

frontend/src/core/codemirror/language/languages/sql/sql.ts

@@ -17,14 +17,52 @@ import {
   DuckDBDialect,
 } from "@marimo-team/codemirror-sql/dialects";
 import type { DataSourceConnection } from "@/core/kernel/messages";
+import { logNever } from "@/utils/assertNever";
+import { Logger } from "@/utils/Logger";
+
+const KNOWN_DIALECTS_ARRAY = [
+  "postgresql",
+  "postgres",
+  "db2",
+  "mysql",
+  "sqlite",
+  "mssql",
+  "sqlserver",
+  "duckdb",
+  "mariadb",
+  "cassandra",
+  "noql",
+  "athena",
+  "bigquery",
+  "hive",
+  "redshift",
+  "snowflake",
+  "flink",
+  "mongodb",
+  "oracle",
+  "oracledb",
+  "timescaledb",
+] as const;
+const KNOWN_DIALECTS: ReadonlySet<string> = new Set(KNOWN_DIALECTS_ARRAY);
+type KnownDialect = (typeof KNOWN_DIALECTS_ARRAY)[number];
+
+export function isKnownDialect(dialect: string): dialect is KnownDialect {
+  return KNOWN_DIALECTS.has(dialect);
+}
 
 /**
  * Guess the CodeMirror SQL dialect from the backend connection dialect.
+ * If unknown, return the standard SQL dialect.
  */
 export function guessDialect(
   connection: Pick<DataSourceConnection, "dialect">,
-): SQLDialect | undefined {
-  switch (connection.dialect) {
+): SQLDialect {
+  const dialect = connection.dialect;
+  if (!isKnownDialect(dialect)) {
+    return ModifiedStandardSQL;
+  }
+
+  switch (dialect) {
     case "postgresql":
     case "postgres":
       return PostgreSQL;
@@ -46,8 +84,21 @@ export function guessDialect(
       return PLSQL;
     case "bigquery":
       return BigQueryDialect;
+    case "timescaledb":
+      return PostgreSQL; // TimescaleDB is a PostgreSQL dialect
+    case "athena":
+    case "db2":
+    case "hive":
+    case "redshift":
+    case "snowflake":
+    case "flink":
+    case "mongodb":
+    case "noql":
+      Logger.debug("Unsupported dialect", { dialect });
+      return ModifiedStandardSQL;
     default:
-      return undefined;
+      logNever(dialect);
+      return ModifiedStandardSQL;
   }
 }
 

frontend/src/core/codemirror/language/languages/sql/utils.ts

@@ -270,6 +270,7 @@ describe("add table list", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
@@ -283,6 +284,7 @@ describe("add table list", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     };
 
     const tableList: DataTable[] = [
@@ -344,6 +346,7 @@ describe("add table list", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "non_existent",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
@@ -407,6 +410,7 @@ describe("add table", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
@@ -420,6 +424,7 @@ describe("add table", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "public",
+      dialect: "sqlite",
     };
 
     const table: DataTable = {
@@ -476,6 +481,7 @@ describe("add table", () => {
       engine: "conn1" as ConnectionName,
       database: "db1",
       schema: "non_existent",
+      dialect: "sqlite",
     });
 
     const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);

frontend/src/core/datasets/__tests__/data-source.test.ts

@@ -47,6 +47,7 @@ export interface SQLTableContext {
   engine: string;
   database: string;
   schema: string;
+  dialect: string;
   defaultSchema?: string | null;
   defaultDatabase?: string | null;
 }