feat: add support for database blocks

core/src/databases/database.rs

@@ -63,7 +63,7 @@ impl Database {
                     .par_iter()
                     .map(|(table, r)| {
                         Ok((
-                            table.table_id().to_string(),
+                            table.name().to_string(),
                             DatabaseSchemaTable::new(table.clone(), TableSchema::from_rows(&r)?),
                         ))
                     })
@@ -95,7 +95,7 @@ impl Database {
                 let rows_by_table = match self.get_rows(store.clone()).await {
                     Ok(rows) => Ok(rows
                         .into_iter()
-                        .map(|(table, rows)| (table.table_id().to_string(), rows))
+                        .map(|(table, rows)| (table.name().to_string(), rows))
                         .collect::<HashMap<_, _>>()),
                     _ => Err(anyhow!("Error retrieving rows from database.")),
                 }?;

front/components/app/NewBlock.tsx

@@ -95,6 +95,12 @@ export default function NewBlock({
       name: "While End",
       description: "Loop over a set of blocks until a condition is met.",
     },
+    {
+      type: "database",
+      typeNames: ["database"],
+      name: "Database",
+      description: "Query a database.",
+    },
     {
       type: "database_schema",
       typeNames: ["database_schema"],

front/components/app/SpecRunView.tsx

@@ -11,6 +11,7 @@ import Chat from "./blocks/Chat";
 import Code from "./blocks/Code";
 import Curl from "./blocks/Curl";
 import Data from "./blocks/Data";
+import Database from "./blocks/Database";
 import DatabaseSchema from "./blocks/DatabaseSchema";
 import DataSource from "./blocks/DataSource";
 import Input from "./blocks/Input";
@@ -356,6 +357,26 @@ export default function SpecRunView({
                 />
               );
 
+            case "database":
+              return (
+                <Database
+                  key={idx}
+                  block={block}
+                  owner={owner}
+                  app={app}
+                  spec={spec}
+                  run={run}
+                  status={status}
+                  running={runRequested || run?.status.run == "running"}
+                  readOnly={readOnly}
+                  onBlockUpdate={(block) => handleSetBlock(idx, block)}
+                  onBlockDelete={() => handleDeleteBlock(idx)}
+                  onBlockUp={() => handleMoveBlockUp(idx)}
+                  onBlockDown={() => handleMoveBlockDown(idx)}
+                  onBlockNew={(blockType) => handleNewBlock(idx, blockType)}
+                />
+              );
+
             default:
               return ((t: never) => (
                 <div key={idx} className="flex flex-row px-4 py-4 text-sm">

front/components/app/blocks/Database.tsx

@@ -0,0 +1,164 @@
+import "@uiw/react-textarea-code-editor/dist.css";
+
+import dynamic from "next/dynamic";
+
+import DataSourcePicker from "@app/components/data_source/DataSourcePicker";
+import DatabasePicker from "@app/components/database/DatabasePicker";
+import { classNames, shallowBlockClone } from "@app/lib/utils";
+import { SpecificationBlockType, SpecificationType } from "@app/types/app";
+import { AppType } from "@app/types/app";
+import { BlockType } from "@app/types/run";
+import { RunType } from "@app/types/run";
+import { WorkspaceType } from "@app/types/user";
+
+import Block from "./Block";
+
+const CodeEditor = dynamic(
+  () => import("@uiw/react-textarea-code-editor").then((mod) => mod.default),
+  { ssr: false }
+);
+
+export default function Database({
+  owner,
+  app,
+  spec,
+  run,
+  block,
+  status,
+  running,
+  readOnly,
+  onBlockUpdate,
+  onBlockDelete,
+  onBlockUp,
+  onBlockDown,
+  onBlockNew,
+}: React.PropsWithChildren<{
+  owner: WorkspaceType;
+  app: AppType;
+  spec: SpecificationType;
+  run: RunType | null;
+  block: SpecificationBlockType;
+  status: any;
+  running: boolean;
+  readOnly: boolean;
+  onBlockUpdate: (block: SpecificationBlockType) => void;
+  onBlockDelete: () => void;
+  onBlockUp: () => void;
+  onBlockDown: () => void;
+  onBlockNew: (blockType: BlockType | "map_reduce" | "while_end") => void;
+}>) {
+  return (
+    <Block
+      owner={owner}
+      app={app}
+      spec={spec}
+      run={run}
+      block={block}
+      status={status}
+      running={running}
+      readOnly={readOnly}
+      canUseCache={false}
+      onBlockUpdate={onBlockUpdate}
+      onBlockDelete={onBlockDelete}
+      onBlockUp={onBlockUp}
+      onBlockDown={onBlockDown}
+      onBlockNew={onBlockNew}
+    >
+      <div className="mx-4 flex w-full flex-col gap-2">
+        <div className="flex flex-col flex-col gap-2 pt-4 xl:flex-row">
+          <div className="flex flex-col xl:flex-row xl:space-x-2">
+            <div className="mr-1 flex flex-initial text-sm font-medium leading-8 text-gray-700">
+              Database:
+            </div>
+            <div className="mr-2 flex flex-initial flex-row items-center space-x-1 text-sm font-medium leading-8 text-gray-700">
+              <DataSourcePicker
+                owner={owner}
+                readOnly={readOnly}
+                currentDataSources={
+                  block.config.database?.data_source_id
+                    ? [
+                        {
+                          data_source_id: block.config.database.data_source_id,
+                          workspace_id: block.config.database.workspace_id,
+                        },
+                      ]
+                    : []
+                }
+                onDataSourcesUpdate={(dataSources) => {
+                  if (dataSources.length === 0) {
+                    return;
+                  }
+                  const ds = dataSources[0];
+                  const b = shallowBlockClone(block);
+                  b.config.database = {
+                    workspace_id: ds.workspace_id,
+                    data_source_id: ds.data_source_id,
+                  };
+                  onBlockUpdate(b);
+                }}
+              />
+            </div>
+          </div>
+          {block.config.database?.data_source_id && (
+            <div className="flex flex-col xl:flex-row xl:space-x-2">
+              <div className="mr-2 flex flex-initial flex-row items-center space-x-1 text-sm font-medium leading-8 text-gray-700">
+                <DatabasePicker
+                  owner={owner}
+                  dataSource={{
+                    workspace_id: block.config.database.workspace_id,
+                    data_source_id: block.config.database.data_source_id,
+                  }}
+                  readOnly={readOnly}
+                  currentDatabaseId={block.config.database?.database_id}
+                  onDatabaseUpdate={(database) => {
+                    const b = shallowBlockClone(block);
+                    b.config.database.database_id = database.database_id;
+                    onBlockUpdate(b);
+                  }}
+                />
+              </div>
+            </div>
+          )}
+        </div>
+        {block.config.database?.database_id && (
+          <div>
+            <div className="mr-1 flex flex-initial text-sm font-medium leading-8 text-gray-700">
+              query:
+            </div>
+            <div className="flex w-full font-normal">
+              <div className="w-full leading-5">
+                <div
+                  className={classNames("border border-slate-100 bg-slate-100")}
+                  style={{
+                    minHeight: "48px",
+                  }}
+                >
+                  <CodeEditor
+                    data-color-mode="light"
+                    readOnly={readOnly}
+                    value={block.spec.query}
+                    language="jinja2"
+                    placeholder=""
+                    onChange={(e) => {
+                      const b = shallowBlockClone(block);
+                      b.spec.query = e.target.value;
+                      onBlockUpdate(b);
+                    }}
+                    padding={3}
+                    style={{
+                      color: "rgb(55 65 81)",
+                      fontSize: 13,
+                      fontFamily:
+                        "ui-monospace, SFMono-Regular, SF Mono, Consolas, Liberation Mono, Menlo, monospace",
+                      backgroundColor: "rgb(241 245 249)",
+                    }}
+                  />
+                </div>
+              </div>
+            </div>
+          </div>
+        )}
+      </div>
+    </Block>
+  );
+}

front/lib/config.ts

@@ -97,6 +97,12 @@ export function extractConfig(spec: SpecificationType): BlockRunConfig {
           database: spec[i].config?.database,
         };
         break;
+      case "database":
+        c[spec[i].name] = {
+          type: "database",
+          database: spec[i].config?.database,
+        };
+        break;
       case "data":
       case "code":
       case "map":

front/lib/specification.ts

@@ -251,6 +251,17 @@ export function addBlock(
         config: {},
       });
       break;
+    case "database":
+      s.splice(idx + 1, 0, {
+        type: "database",
+        name: getNextName(spec, "DATABASE"),
+        indent: 0,
+        spec: {
+          query: "",
+        },
+        config: {},
+      });
+      break;
     default:
       s.splice(idx + 1, 0, {
         type: blockType,
@@ -568,6 +579,15 @@ export function dumpSpecification(
         out += "\n";
         break;
       }
+      case "database": {
+        out += `database ${block.name} {\n`;
+        out += `  query: \n\`\`\`\n${escapeTripleBackticks(
+          block.spec.query
+        )}\n\`\`\`\n`;
+        out += `}\n`;
+        out += "\n";
+        break;
+      }
       default:
         ((t: never) => {
           console.error(`Unknown block type: ${t}`);

front/types/run.ts

@@ -14,7 +14,8 @@ export type BlockType =
   | "search"
   | "curl"
   | "browser"
-  | "database_schema";
+  | "database_schema"
+  | "database";
 
 export type RunRunType = "deploy" | "local" | "execute";
 type Status = "running" | "succeeded" | "errored";