From 1934f72cbb9b24cdb6af78471c1b95fa99a0d2ec Mon Sep 17 00:00:00 2001 From: Henry Fontanier Date: Wed, 22 Nov 2023 17:49:07 +0100 Subject: [PATCH] feat: database_schema support in dust apps (#2624) * Add database block * add io-ts-types * add front support for database schema block * rm debug log * improve copy * nits * extra nit --- core/src/blocks/block.rs | 3 + front/components/app/NewBlock.tsx | 22 +-- front/components/app/SpecRunView.tsx | 39 ++++-- .../components/app/blocks/DatabaseSchema.tsx | 117 ++++++++++++++++ front/components/database/DatabasePicker.tsx | 115 ++++++++++++++++ front/lib/config.ts | 22 ++- front/lib/specification.ts | 18 +++ front/lib/swr.ts | 27 ++++ front/package-lock.json | 31 +++++ front/package.json | 1 + .../data_sources/[name]/databases/index.ts | 13 +- .../data_sources/[name]/databases/index.ts | 129 ++++++++++++++++++ front/types/run.ts | 3 +- 13 files changed, 509 insertions(+), 31 deletions(-) create mode 100644 front/components/app/blocks/DatabaseSchema.tsx create mode 100644 front/components/database/DatabasePicker.tsx create mode 100644 front/pages/api/w/[wId]/data_sources/[name]/databases/index.ts diff --git a/core/src/blocks/block.rs b/core/src/blocks/block.rs index e98389ea36b1..d4dcba5acf96 100644 --- a/core/src/blocks/block.rs +++ b/core/src/blocks/block.rs @@ -74,6 +74,7 @@ pub enum BlockType { Browser, While, End, + #[serde(rename = "database_schema")] DatabaseSchema, Database, } @@ -117,6 +118,8 @@ impl FromStr for BlockType { "browser" => Ok(BlockType::Browser), "while" => Ok(BlockType::While), "end" => Ok(BlockType::End), + "database_schema" => Ok(BlockType::DatabaseSchema), + "database" => Ok(BlockType::Database), _ => Err(ParseError::with_message("Unknown BlockType"))?, } } diff --git a/front/components/app/NewBlock.tsx b/front/components/app/NewBlock.tsx index 2599c7bb62df..2fa21d8f0993 100644 --- a/front/components/app/NewBlock.tsx +++ b/front/components/app/NewBlock.tsx @@ -21,7 +21,12 @@ export default function NewBlock({ }) { const containsInput = spec.filter((block) => block.type == "input").length > 0; - const blocks = [ + const blocks: { + type: BlockType | "map_reduce" | "while_end"; + typeNames: BlockType[]; + name: string; + description: string; + }[] = [ { type: "llm", typeNames: ["llm"], @@ -90,12 +95,13 @@ export default function NewBlock({ name: "While End", description: "Loop over a set of blocks until a condition is met.", }, - ] as { - type: BlockType | "map_reduce" | "while_end"; - typeNames: BlockType[]; - name: string; - description: string; - }[]; + { + type: "database_schema", + typeNames: ["database_schema"], + name: "Database Schema", + description: "Retrieve the schema of a database.", + }, + ]; if (!containsInput) { blocks.splice(0, 0, { @@ -164,7 +170,7 @@ export default function NewBlock({ ))} -
+
{block.description}
diff --git a/front/components/app/SpecRunView.tsx b/front/components/app/SpecRunView.tsx index fca946f1bcf2..fd002e36976a 100644 --- a/front/components/app/SpecRunView.tsx +++ b/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 DatabaseSchema from "./blocks/DatabaseSchema"; import DataSource from "./blocks/DataSource"; import Input from "./blocks/Input"; import LLM from "./blocks/LLM"; @@ -94,7 +95,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "data": return ( @@ -115,7 +115,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "llm": return ( @@ -136,7 +135,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "chat": return ( @@ -157,7 +155,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "code": return ( @@ -178,7 +175,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "data_source": return ( @@ -199,7 +195,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "map": return ( @@ -220,7 +215,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "reduce": return ( @@ -261,7 +255,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "end": return ( @@ -302,7 +295,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "curl": return ( @@ -323,7 +315,6 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; case "browser": return ( @@ -344,15 +335,33 @@ export default function SpecRunView({ onBlockNew={(blockType) => handleNewBlock(idx, blockType)} /> ); - break; - default: + case "database_schema": return ( + handleSetBlock(idx, block)} + onBlockDelete={() => handleDeleteBlock(idx)} + onBlockUp={() => handleMoveBlockUp(idx)} + onBlockDown={() => handleMoveBlockDown(idx)} + onBlockNew={(blockType) => handleNewBlock(idx, blockType)} + /> + ); + + default: + return ((t: never) => (
- Unknown block type: {block.type} + Unknown block type: {t}
- ); - break; + ))(block.type); } })} diff --git a/front/components/app/blocks/DatabaseSchema.tsx b/front/components/app/blocks/DatabaseSchema.tsx new file mode 100644 index 000000000000..812e70cf1c23 --- /dev/null +++ b/front/components/app/blocks/DatabaseSchema.tsx @@ -0,0 +1,117 @@ +import "@uiw/react-textarea-code-editor/dist.css"; + +import DataSourcePicker from "@app/components/data_source/DataSourcePicker"; +import DatabasePicker from "@app/components/database/DatabasePicker"; +import { 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"; + +export default function DatabaseSchema({ + 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 ( + +
+
+
+ Database: +
+
+ { + 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); + }} + /> +
+
+ {block.config.database?.data_source_id && ( +
+
+ { + const b = shallowBlockClone(block); + b.config.database.database_id = database.database_id; + onBlockUpdate(b); + }} + /> +
+
+ )} +
+
+ ); +} diff --git a/front/components/database/DatabasePicker.tsx b/front/components/database/DatabasePicker.tsx new file mode 100644 index 000000000000..b0f3f71d274d --- /dev/null +++ b/front/components/database/DatabasePicker.tsx @@ -0,0 +1,115 @@ +import { Menu } from "@headlessui/react"; +import { ChevronDownIcon } from "@heroicons/react/20/solid"; + +import { CoreAPIDatabase } from "@app/lib/core_api"; +import { useDatabases } from "@app/lib/swr"; +import { classNames } from "@app/lib/utils"; +import { WorkspaceType } from "@app/types/user"; + +export default function DatabasePicker({ + owner, + dataSource, + currentDatabaseId, + readOnly, + onDatabaseUpdate, +}: { + owner: WorkspaceType; + dataSource: { + workspace_id: string; + data_source_id: string; + }; + currentDatabaseId?: string; + readOnly: boolean; + onDatabaseUpdate: (database: CoreAPIDatabase) => void; +}) { + void owner; + void dataSource; + + const { databases } = useDatabases({ + workspaceId: dataSource.workspace_id, + dataSourceName: dataSource.data_source_id, + offset: 0, + // TODO(@fontanierh): + // will break if we have more than 100 databases (but UI wouldn't work anyway) + limit: 100, + }); + + const currentDatabase = currentDatabaseId + ? databases.find((db) => db.database_id === currentDatabaseId) + : null; + + return ( +
+
+ {readOnly ? ( + currentDatabase ? ( +
+ {currentDatabase.name} +
+ ) : ( + "No Database" + ) + ) : ( + +
+ + {currentDatabase ? ( + <> +
+ {currentDatabase.name} +
+ + + ) : databases && databases.length > 0 ? ( + "Select Database" + ) : ( + "No Databases" + )} +
+
+ + {(databases || []).length > 0 ? ( + +
+ {(databases || []).map((db) => { + return ( + + {({ active }) => ( + onDatabaseUpdate(db)} + > + {db.name} + + )} + + ); + })} +
+
+ ) : null} +
+ )} +
+
+ ); +} diff --git a/front/lib/config.ts b/front/lib/config.ts index dcdb9bb59d93..0e02d7773f00 100644 --- a/front/lib/config.ts +++ b/front/lib/config.ts @@ -3,7 +3,8 @@ import { BlockRunConfig, SpecificationType } from "@app/types/app"; export function extractConfig(spec: SpecificationType): BlockRunConfig { const c = {} as { [key: string]: any }; for (let i = 0; i < spec.length; i++) { - switch (spec[i].type) { + const type = spec[i].type; + switch (type) { case "llm": c[spec[i].name] = { type: "llm", @@ -90,6 +91,25 @@ export function extractConfig(spec: SpecificationType): BlockRunConfig { : false, }; break; + case "database_schema": + c[spec[i].name] = { + type: "database_schema", + database: spec[i].config?.database, + }; + break; + case "data": + case "code": + case "map": + case "reduce": + case "while": + case "end": + // these blocks have no config + break; + + default: + ((t: never) => { + console.warn(`Unknown block type: ${t}`); + })(type); } } return c; diff --git a/front/lib/specification.ts b/front/lib/specification.ts index de8ca7591086..81086a95363c 100644 --- a/front/lib/specification.ts +++ b/front/lib/specification.ts @@ -242,6 +242,15 @@ export function addBlock( }, }); break; + case "database_schema": + s.splice(idx + 1, 0, { + type: "database_schema", + name: getNextName(spec, "DATABASE_SCHEMA"), + indent: 0, + spec: {}, + config: {}, + }); + break; default: s.splice(idx + 1, 0, { type: blockType, @@ -554,6 +563,15 @@ export function dumpSpecification( out += "\n"; break; } + case "database_schema": { + out += `database_schema ${block.name} { }\n`; + out += "\n"; + break; + } + default: + ((t: never) => { + console.error(`Unknown block type: ${t}`); + })(block.type); } } return out.slice(0, -1); diff --git a/front/lib/swr.ts b/front/lib/swr.ts index 3df5f6c83833..f5f05bf9da00 100644 --- a/front/lib/swr.ts +++ b/front/lib/swr.ts @@ -3,6 +3,7 @@ import useSWR, { Fetcher } from "swr"; import { GetPokePlansResponseBody } from "@app/pages/api/poke/plans"; import { GetWorkspacesResponseBody } from "@app/pages/api/poke/workspaces"; import { GetUserMetadataResponseBody } from "@app/pages/api/user/metadata/[key]"; +import { ListDatabasesResponseBody } from "@app/pages/api/v1/w/[wId]/data_sources/[name]/databases"; import { GetDatasetsResponseBody } from "@app/pages/api/w/[wId]/apps/[aId]/datasets"; import { GetRunsResponseBody } from "@app/pages/api/w/[wId]/apps/[aId]/runs"; import { GetRunBlockResponseBody } from "@app/pages/api/w/[wId]/apps/[aId]/runs/[runId]/blocks/[type]/[name]"; @@ -504,3 +505,29 @@ export function useSlackChannelsLinkedWithAgent({ mutateSlackChannels: mutate, }; } + +export function useDatabases({ + workspaceId, + dataSourceName, + offset, + limit, +}: { + workspaceId: string; + dataSourceName: string; + offset: number; + limit: number; +}) { + const databasesFetcher: Fetcher = fetcher; + + const { data, error, mutate } = useSWR( + `/api/w/${workspaceId}/data_sources/${dataSourceName}/databases?offset=${offset}&limit=${limit}`, + databasesFetcher + ); + + return { + databases: data ? data.databases : [], + isDatabasesLoading: !error && !data, + isDatabasesError: error, + mutateDatabases: mutate, + }; +} diff --git a/front/package-lock.json b/front/package-lock.json index 4d3fde8736f0..3a9f3ce94030 100644 --- a/front/package-lock.json +++ b/front/package-lock.json @@ -34,6 +34,7 @@ "hot-shots": "^10.0.0", "io-ts": "^2.2.20", "io-ts-reporters": "^2.0.1", + "io-ts-types": "^0.5.19", "jsonwebtoken": "^9.0.0", "minimist": "^1.2.8", "moment-timezone": "^0.5.43", @@ -7417,6 +7418,17 @@ "io-ts": "^2.2.16" } }, + "node_modules/io-ts-types": { + "version": "0.5.19", + "resolved": "https://registry.npmjs.org/io-ts-types/-/io-ts-types-0.5.19.tgz", + "integrity": "sha512-kQOYYDZG5vKre+INIDZbLeDJe+oM+4zLpUkjXyTMyUfoCpjJNyi29ZLkuEAwcPufaYo3yu/BsemZtbdD+NtRfQ==", + "peerDependencies": { + "fp-ts": "^2.0.0", + "io-ts": "^2.0.0", + "monocle-ts": "^2.0.0", + "newtype-ts": "^0.3.2" + } + }, "node_modules/ip": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/ip/-/ip-2.0.0.tgz", @@ -10224,6 +10236,15 @@ "node": "*" } }, + "node_modules/monocle-ts": { + "version": "2.3.13", + "resolved": "https://registry.npmjs.org/monocle-ts/-/monocle-ts-2.3.13.tgz", + "integrity": "sha512-D5Ygd3oulEoAm3KuGO0eeJIrhFf1jlQIoEVV2DYsZUMz42j4tGxgct97Aq68+F8w4w4geEnwFa8HayTS/7lpKQ==", + "peer": true, + "peerDependencies": { + "fp-ts": "^2.5.0" + } + }, "node_modules/mri": { "version": "1.2.0", "resolved": "https://registry.npmjs.org/mri/-/mri-1.2.0.tgz", @@ -10313,6 +10334,16 @@ "resolved": "https://registry.npmjs.org/neo-async/-/neo-async-2.6.2.tgz", "integrity": "sha512-Yd3UES5mWCSqR+qNT93S3UoYUkqAZ9lLg8a7g9rimsWmYGK8cVToA4/sF3RrshdyV3sAGMXVUmpMYOw+dLpOuw==" }, + "node_modules/newtype-ts": { + "version": "0.3.5", + "resolved": "https://registry.npmjs.org/newtype-ts/-/newtype-ts-0.3.5.tgz", + "integrity": "sha512-v83UEQMlVR75yf1OUdoSFssjitxzjZlqBAjiGQ4WJaML8Jdc68LJ+BaSAXUmKY4bNzp7hygkKLYTsDi14PxI2g==", + "peer": true, + "peerDependencies": { + "fp-ts": "^2.0.0", + "monocle-ts": "^2.0.0" + } + }, "node_modules/next": { "version": "13.5.4", "resolved": "https://registry.npmjs.org/next/-/next-13.5.4.tgz", diff --git a/front/package.json b/front/package.json index 78cef2513c8d..6a409faecd64 100644 --- a/front/package.json +++ b/front/package.json @@ -42,6 +42,7 @@ "hot-shots": "^10.0.0", "io-ts": "^2.2.20", "io-ts-reporters": "^2.0.1", + "io-ts-types": "^0.5.19", "jsonwebtoken": "^9.0.0", "minimist": "^1.2.8", "moment-timezone": "^0.5.43", diff --git a/front/pages/api/v1/w/[wId]/data_sources/[name]/databases/index.ts b/front/pages/api/v1/w/[wId]/data_sources/[name]/databases/index.ts index eb8fd9a0a6e9..7db72260b97b 100644 --- a/front/pages/api/v1/w/[wId]/data_sources/[name]/databases/index.ts +++ b/front/pages/api/v1/w/[wId]/data_sources/[name]/databases/index.ts @@ -1,6 +1,7 @@ import { isLeft } from "fp-ts/lib/Either"; import * as t from "io-ts"; import * as reporter from "io-ts-reporters"; +import { IntFromString } from "io-ts-types"; import { NextApiRequest, NextApiResponse } from "next"; import { getDataSource } from "@app/lib/api/data_sources"; @@ -11,18 +12,18 @@ import { generateModelSId } from "@app/lib/utils"; import logger from "@app/logger/logger"; import { apiError, withLogging } from "@app/logger/withlogging"; -const CreateDatabaseReqBodySchema = t.type({ +export const CreateDatabaseReqBodySchema = t.type({ name: t.string, }); -type CreateDatabaseResponseBody = { +export type CreateDatabaseResponseBody = { database: CoreAPIDatabase; }; -const ListDatabasesReqQuerySchema = t.type({ - offset: t.number, - limit: t.number, +export const ListDatabasesReqQuerySchema = t.type({ + offset: t.union([IntFromString, t.number]), + limit: t.union([IntFromString, t.number]), }); -type ListDatabasesResponseBody = { +export type ListDatabasesResponseBody = { databases: CoreAPIDatabase[]; }; diff --git a/front/pages/api/w/[wId]/data_sources/[name]/databases/index.ts b/front/pages/api/w/[wId]/data_sources/[name]/databases/index.ts new file mode 100644 index 000000000000..07daed0d7dd9 --- /dev/null +++ b/front/pages/api/w/[wId]/data_sources/[name]/databases/index.ts @@ -0,0 +1,129 @@ +import { isLeft } from "fp-ts/lib/Either"; +import * as reporter from "io-ts-reporters"; +import { NextApiRequest, NextApiResponse } from "next"; + +import { getDataSource } from "@app/lib/api/data_sources"; +import { Authenticator, getSession } from "@app/lib/auth"; +import { CoreAPI } from "@app/lib/core_api"; +import { isDevelopmentOrDustWorkspace } from "@app/lib/development"; +import logger from "@app/logger/logger"; +import { apiError, withLogging } from "@app/logger/withlogging"; +import { + ListDatabasesReqQuerySchema, + ListDatabasesResponseBody, +} from "@app/pages/api/v1/w/[wId]/data_sources/[name]/databases"; + +async function handler( + req: NextApiRequest, + res: NextApiResponse +): Promise { + const session = await getSession(req, res); + const auth = await Authenticator.fromSession( + session, + req.query.wId as string + ); + + const owner = auth.workspace(); + if (!owner || !auth.isBuilder()) { + return apiError(req, res, { + status_code: 404, + api_error: { + type: "data_source_not_found", + message: "The data source you requested was not found.", + }, + }); + } + + const plan = auth.plan(); + if (!owner || !plan) { + return apiError(req, res, { + status_code: 404, + api_error: { + type: "workspace_not_found", + message: "The workspace you requested was not found.", + }, + }); + } + + if (!isDevelopmentOrDustWorkspace(owner)) { + res.status(404).end(); + return; + } + + if (!req.query.name || typeof req.query.name !== "string") { + return apiError(req, res, { + status_code: 404, + api_error: { + type: "data_source_not_found", + message: "The data source you requested was not found.", + }, + }); + } + const dataSource = await getDataSource(auth, req.query.name as string); + if (!dataSource) { + return apiError(req, res, { + status_code: 404, + api_error: { + type: "data_source_not_found", + message: "The data source you requested was not found.", + }, + }); + } + + switch (req.method) { + case "GET": + const queryValidation = ListDatabasesReqQuerySchema.decode(req.query); + if (isLeft(queryValidation)) { + const pathError = reporter.formatValidationErrors(queryValidation.left); + return apiError(req, res, { + api_error: { + type: "invalid_request_error", + message: `Invalid request query: ${pathError}`, + }, + status_code: 400, + }); + } + + const { offset, limit } = queryValidation.right; + + const getRes = await CoreAPI.getDatabases({ + projectId: dataSource.dustAPIProjectId, + dataSourceName: dataSource.name, + offset, + limit, + }); + + if (getRes.isErr()) { + logger.error( + { + dataSourceName: dataSource.name, + workspaceId: owner.id, + error: getRes.error, + }, + "Failed to list databases." + ); + return apiError(req, res, { + status_code: 500, + api_error: { + type: "internal_server_error", + message: "Failed to list databases.", + }, + }); + } + + const { databases } = getRes.value; + + return res.status(200).json({ databases }); + + default: + return apiError(req, res, { + status_code: 405, + api_error: { + type: "method_not_supported_error", + message: "The method passed is not supported, GET is expected.", + }, + }); + } +} + +export default withLogging(handler); diff --git a/front/types/run.ts b/front/types/run.ts index 7707cf9dd1a7..3eabf91aab40 100644 --- a/front/types/run.ts +++ b/front/types/run.ts @@ -13,7 +13,8 @@ export type BlockType = | "end" | "search" | "curl" - | "browser"; + | "browser" + | "database_schema"; export type RunRunType = "deploy" | "local" | "execute"; type Status = "running" | "succeeded" | "errored";