generate different sql code based on dialect

frontend/src/components/datasources/datasources.tsx

@@ -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/utils.ts

@@ -1,4 +1,7 @@
 /* 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";
@@ -9,6 +12,59 @@ 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,
+      };
+    default:
+      return defaultFormatter;
+  }
+}
+
 export function sqlCode({
   table,
   columnName,
@@ -19,8 +75,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 +101,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"""${selectClause}""")`;
     }
 
-    return `_df = mo.sql(f"SELECT ${columnName} FROM ${tableName} LIMIT 100", engine=${engine})`;
+    return `_df = mo.sql(f"""${selectClause}""", engine=${engine})`;
   }
 
   return `_df = mo.sql(f'SELECT "${columnName}" FROM ${table.name} LIMIT 100')`;

frontend/src/core/codemirror/language/languages/sql/completion-store.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/sql.ts

@@ -52,6 +52,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 +354,11 @@ function connectionNameToParserDialect(
 ): ParserDialects | null {
   const dialect =
     SCHEMA_CACHE.getInternalDialect(connectionName)?.toLowerCase();
+
+  if (!dialect || !isKnownDialect(dialect)) {
+    return null;
+  }
+
   switch (dialect) {
     case "postgresql":
     case "postgres":

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

@@ -18,13 +18,49 @@ import {
 } from "@marimo-team/codemirror-sql/dialects";
 import type { DataSourceConnection } from "@/core/kernel/messages";
 
+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;
@@ -47,7 +83,7 @@ export function guessDialect(
     case "bigquery":
       return BigQueryDialect;
     default:
-      return undefined;
+      return ModifiedStandardSQL;
   }
 }
 

frontend/src/core/datasets/__tests__/data-source.test.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/data-source-connections.ts

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