@marimo-team/islands 0.17.2 → 0.17.3

package/src/components/data-table/column-summary/legacy-chart-spec.ts

@@ -1,10 +1,22 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-// New spec is under a feature flag, we keep this until full migration
+// Fallback spec
 
 import { mint, orange, slate } from "@radix-ui/colors";
-import type { Scale } from "vega-lite/build/src/scale";
+import type { TopLevelSpec } from "vega-lite";
+import type { Scale } from "vega-lite/types_unstable/scale.js";
+// @ts-expect-error vega-typings does not include formats
+import { formats } from "vega-loader";
+import { asRemoteURL } from "@/core/runtime/config";
 import type { TopLevelFacetedUnitSpec } from "@/plugins/impl/data-explorer/queries/types";
+import { arrow } from "@/plugins/impl/vega/formats";
+import { parseCsvData } from "@/plugins/impl/vega/loader";
+import {
+  byteStringToBinary,
+  extractBase64FromDataURL,
+  isDataURLString,
+  typedAtob,
+} from "@/utils/json/base64";
 
 export function getLegacyNumericSpec(
   column: string,
@@ -240,3 +252,67 @@ export function getLegacyBooleanSpec(
     ],
   } as TopLevelFacetedUnitSpec; // "layer" not in TopLevelFacetedUnitSpec
 }
+
+// Arrow formats have a magic number at the beginning of the file.
+const ARROW_MAGIC_NUMBER = "ARROW1";
+
+// register arrow reader under type 'arrow'
+formats("arrow", arrow);
+
+export function getDataSpecAndSourceName<T>(data: string | T[]): {
+  dataSpec: TopLevelSpec["data"];
+  sourceName: "data_0" | "source_0";
+} {
+  let dataSpec: TopLevelSpec["data"];
+  let sourceName: "data_0" | "source_0";
+
+  // Data may come in from a few different sources:
+  // - A URL
+  // - A CSV data URI (e.g. "data:text/csv;base64,...")
+  // - A CSV string (e.g. "a,b,c\n1,2,3\n4,5,6")
+  // - An array of objects
+  // For each case, we need to set up the data spec and source name appropriately.
+  // If its a file, the source name will be "source_0", otherwise it will be "data_0".
+  // We have a few snapshot tests to ensure that the spec is correct for each case.
+  if (typeof data === "string") {
+    if (data.startsWith("./@file") || data.startsWith("/@file")) {
+      dataSpec = { url: asRemoteURL(data).href };
+      sourceName = "source_0";
+    } else if (isDataURLString(data)) {
+      sourceName = "data_0";
+      const base64 = extractBase64FromDataURL(data);
+      const decoded = typedAtob(base64);
+
+      // eslint-disable-next-line unicorn/prefer-ternary
+      if (decoded.startsWith(ARROW_MAGIC_NUMBER)) {
+        // @ts-expect-error vega-typings does not include arrow format
+        dataSpec = {
+          values: byteStringToBinary(decoded),
+          format: { type: "arrow" },
+        };
+      } else {
+        // Assume it's a CSV string
+        dataSpec = { values: parseCsvData(decoded) };
+      }
+    } else {
+      // Assume it's a CSV string
+      dataSpec = { values: parseCsvData(data) };
+      sourceName = "data_0";
+    }
+  } else {
+    dataSpec = { values: data };
+    sourceName = "source_0";
+  }
+
+  return { dataSpec, sourceName };
+}
+
+export function getScale(sourceName: string): Scale {
+  return {
+    align: 0,
+    paddingInner: 0,
+    paddingOuter: {
+      expr: `length(data('${sourceName}')) == 2 ? 1 : length(data('${sourceName}')) == 3 ? 0.5 : length(data('${sourceName}')) == 4 ? 0 : 0`,
+    },
+  };
+}

package/src/components/data-table/column-summary/utils.ts

@@ -1,6 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { StringFieldDef } from "vega-lite/build/src/channeldef";
+import type { StringFieldDef } from "vega-lite/types_unstable/channeldef.js";
 import { Logger } from "@/utils/Logger";
 import type { BinValues } from "../types";
 

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

@@ -24,154 +24,355 @@ describe("sqlCode", () => {
     name: "email" as const,
   } as DataTableColumn;
 
-  it("should generate basic SQL without sqlTableContext", () => {
-    const result = sqlCode({ table: mockTable, columnName: mockColumn.name });
-    expect(result).toBe(
-      "_df = mo.sql(f'SELECT \"email\" FROM users LIMIT 100')",
-    );
-  });
+  describe("basic SQL generation", () => {
+    it("should generate basic SQL without sqlTableContext", () => {
+      const result = sqlCode({ table: mockTable, columnName: mockColumn.name });
+      expect(result).toBe(
+        "_df = mo.sql(f'SELECT \"email\" FROM users LIMIT 100')",
+      );
+    });
 
-  it("should generate SQL with default schema", () => {
-    const sqlTableContext: SQLTableContext = {
-      engine: DUCKDB_ENGINE,
-      schema: "public",
-      defaultSchema: "public",
-      defaultDatabase: "mydb",
-      database: "mydb",
-    };
-
-    const result = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext,
-    });
-    expect(result).toBe('_df = mo.sql(f"SELECT email FROM users LIMIT 100")');
-  });
+    it("should generate SQL with default schema", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: DUCKDB_ENGINE,
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "duckdb",
+      };
 
-  it("should generate SQL with non-default schema", () => {
-    const sqlTableContext: SQLTableContext = {
-      engine: DUCKDB_ENGINE,
-      schema: "analytics",
-      defaultSchema: "public",
-      defaultDatabase: "mydb",
-      database: "mydb",
-    };
-
-    const result = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext,
-    });
-    expect(result).toBe(
-      '_df = mo.sql(f"SELECT email FROM analytics.users LIMIT 100")',
-    );
-  });
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "users" LIMIT 100\n""")',
+      );
+    });
+
+    it("should generate SQL with non-default schema", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: DUCKDB_ENGINE,
+        schema: "analytics",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "duckdb",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "analytics"."users" LIMIT 100\n""")',
+      );
+    });
+
+    it("should generate SQL with non-default engine", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "snowflake",
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "snowflake",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT email FROM users LIMIT 100\n""", engine=snowflake)',
+      );
+    });
+
+    it("should generate SQL with non-default database", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: DUCKDB_ENGINE,
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "memory",
+        database: "remote",
+        dialect: "duckdb",
+      };
 
-  it("should generate SQL with non-default engine", () => {
-    const sqlTableContext: SQLTableContext = {
-      engine: "snowflake",
-      schema: "public",
-      defaultSchema: "public",
-      defaultDatabase: "mydb",
-      database: "mydb",
-    };
-
-    const result = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext,
-    });
-    expect(result).toBe(
-      '_df = mo.sql(f"SELECT email FROM users LIMIT 100", engine=snowflake)',
-    );
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "remote"."users" LIMIT 100\n""")',
+      );
+    });
+
+    it("should generate SQL for schemaless tables", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: DUCKDB_ENGINE,
+        schema: "",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "duckdb",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "users" LIMIT 100\n""")',
+      );
+
+      const sqlTableContext2: SQLTableContext = {
+        engine: DUCKDB_ENGINE,
+        schema: "",
+        defaultDatabase: "remote",
+        database: "another_db",
+        dialect: "duckdb",
+      };
+
+      const result2 = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext: sqlTableContext2,
+      });
+      expect(result2).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "another_db"."users" LIMIT 100\n""")',
+      );
+    });
   });
 
-  it("should generate SQL with non-default schema and non-default engine", () => {
-    const sqlTableContext: SQLTableContext = {
-      engine: "bigquery",
-      schema: "sales",
-      defaultSchema: "public",
-      defaultDatabase: "mydb",
-      database: "mydb",
-    };
-
-    const result = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext,
-    });
-    expect(result).toBe(
-      '_df = mo.sql(f"SELECT email FROM sales.users LIMIT 100", engine=bigquery)',
-    );
+  describe("BigQuery dialect", () => {
+    it("should use backticks for table names", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "bigquery",
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "bigquery",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT email FROM `users` LIMIT 100\n""", engine=bigquery)',
+      );
+    });
+
+    it("should use backticks for database.schema.table", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "bigquery",
+        schema: "sales",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "remote",
+        dialect: "bigquery",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT email FROM `remote.sales.users` LIMIT 100\n""", engine=bigquery)',
+      );
+    });
+
+    it("should handle case-insensitive dialect name", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "bigquery",
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "BigQuery",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT email FROM `users` LIMIT 100\n""", engine=bigquery)',
+      );
+    });
   });
 
-  it("should generate SQL with non-default database", () => {
-    const sqlTableContext: SQLTableContext = {
-      engine: DUCKDB_ENGINE,
-      schema: "public",
-      defaultSchema: "public",
-      defaultDatabase: "memory",
-      database: "remote",
-    };
-
-    const result = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext,
-    });
-    expect(result).toBe(
-      '_df = mo.sql(f"SELECT email FROM remote.users LIMIT 100")',
-    );
+  describe("MSSQL dialect", () => {
+    it("should use TOP 100 instead of LIMIT", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "mssql",
+        schema: "dbo",
+        defaultSchema: "dbo",
+        defaultDatabase: "master",
+        database: "master",
+        dialect: "mssql",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT TOP 100 email FROM users\n""", engine=mssql)',
+      );
+    });
+
+    it("should use TOP 100 with database and schema prefix", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "mssql",
+        schema: "sales",
+        defaultSchema: "dbo",
+        defaultDatabase: "master",
+        database: "analytics",
+        dialect: "mssql",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT TOP 100 email FROM analytics.sales.users\n""", engine=mssql)',
+      );
+    });
   });
 
-  it("should generate SQL with non-default database and non-default schema", () => {
-    const sqlTableContext: SQLTableContext = {
-      engine: "bigquery",
-      schema: "sales",
-      defaultSchema: "public",
-      defaultDatabase: "mydb",
-      database: "remote",
-    };
-
-    const result = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext,
-    });
-    expect(result).toBe(
-      '_df = mo.sql(f"SELECT email FROM remote.sales.users LIMIT 100", engine=bigquery)',
-    );
+  describe("TimescaleDB and PostgreSQL dialect", () => {
+    it("should wrap table name with double quotes", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "timescaledb",
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "timescaledb",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "users" LIMIT 100\n""", engine=timescaledb)',
+      );
+    });
+
+    it("should wrap database, schema, and table name with double quotes", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "timescaledb",
+        schema: "sales",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "remote",
+        dialect: "timescaledb",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "remote"."sales"."users" LIMIT 100\n""", engine=timescaledb)',
+      );
+    });
+
+    it("should handle schemaless database with double quotes", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "postgres",
+        schema: "",
+        defaultDatabase: "mydb",
+        database: "remote",
+        dialect: "postgres",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT "email" FROM "remote"."users" LIMIT 100\n""", engine=postgres)',
+      );
+    });
+
+    it("should not quote * column name", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "postgres",
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "postgres",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: "*",
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT * FROM "users" LIMIT 100\n""", engine=postgres)',
+      );
+    });
   });
 
-  it("should generate SQL for schemaless tables", () => {
-    const sqlTableContext: SQLTableContext = {
-      engine: DUCKDB_ENGINE,
-      schema: "",
-      defaultDatabase: "mydb",
-      database: "mydb",
-    };
-
-    const result = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext,
-    });
-    expect(result).toBe('_df = mo.sql(f"SELECT email FROM users LIMIT 100")');
-
-    const sqlTableContext2: SQLTableContext = {
-      engine: DUCKDB_ENGINE,
-      schema: "",
-      defaultDatabase: "remote",
-      database: "another_db",
-    };
-
-    const result2 = sqlCode({
-      table: mockTable,
-      columnName: mockColumn.name,
-      sqlTableContext: sqlTableContext2,
-    });
-    expect(result2).toBe(
-      '_df = mo.sql(f"SELECT email FROM another_db.users LIMIT 100")',
-    );
+  describe("fallback behavior", () => {
+    it("should use default formatter for unknown dialect", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "postgres",
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "unknown_dialect",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT email FROM users LIMIT 100\n""", engine=postgres)',
+      );
+    });
+
+    it("should use default formatter when dialect is not provided", () => {
+      const sqlTableContext: SQLTableContext = {
+        engine: "postgres",
+        schema: "public",
+        defaultSchema: "public",
+        defaultDatabase: "mydb",
+        database: "mydb",
+        dialect: "",
+      };
+
+      const result = sqlCode({
+        table: mockTable,
+        columnName: mockColumn.name,
+        sqlTableContext,
+      });
+      expect(result).toBe(
+        '_df = mo.sql(f"""\nSELECT email FROM users LIMIT 100\n""", engine=postgres)',
+      );
+    });
   });
 });

package/src/components/datasources/column-preview.tsx

@@ -21,6 +21,7 @@ import type { TopLevelFacetedUnitSpec } from "@/plugins/impl/data-explorer/queri
 import { type Theme, useTheme } from "@/theme/useTheme";
 import { Events } from "@/utils/events";
 import { prettyNumber } from "@/utils/numbers";
+import { LazyVegaEmbed } from "../charts/lazy";
 import type { ColumnHeaderStatsKey } from "../data-table/types";
 import { CopyClipboardIcon } from "../icons/copy-icon";
 import { Spinner } from "../icons/spinner";
@@ -30,10 +31,6 @@ import { ColumnPreviewContainer } from "./components";
 import { InstallPackageButton } from "./install-package-button";
 import { convertStatsName, sqlCode } from "./utils";
 
-const LazyVegaLite = React.lazy(() =>
-  import("react-vega").then((m) => ({ default: m.VegaLite })),
-);
-
 export const DatasetColumnPreview: React.FC<{
   table: DataTable;
   column: DataTableColumn;
@@ -226,12 +223,15 @@ export function renderChart(chartSpec: string, theme: Theme) {
 
   return (
     <Suspense fallback={LoadingChart}>
-      <LazyVegaLite
+      <LazyVegaEmbed
         spec={updateSpec(JSON.parse(chartSpec) as TopLevelFacetedUnitSpec)}
-        width={"container" as unknown as number}
-        height={100}
-        actions={false}
-        theme={theme === "dark" ? "dark" : "vox"}
+        options={{
+          theme: theme === "dark" ? "dark" : "vox",
+          height: 100,
+          width: "container" as unknown as number,
+          actions: false,
+          renderer: "canvas",
+        }}
       />
     </Suspense>
   );

package/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>