diff --git a/src/app/(outerbase)/local-setting/page.tsx b/src/app/(outerbase)/local-setting/page.tsx new file mode 100644 index 00000000..16602022 --- /dev/null +++ b/src/app/(outerbase)/local-setting/page.tsx @@ -0,0 +1,63 @@ +"use client"; +import LabelInput from "@/components/label-input"; +import { Button } from "@/components/orbit/button"; +import { + getAgentFromLocalStorage, + updateAgentFromLocalStorage, +} from "@/lib/ai-agent-storage"; +import { useCallback, useEffect, useState } from "react"; +import { toast } from "sonner"; +import NavigationLayout from "../nav-layout"; + +export default function LocalSettingPage() { + const [token, setToken] = useState(""); + + useEffect(() => { + if (typeof window === "undefined") return; + + const agentData = getAgentFromLocalStorage(); + if (!agentData) return; + + setToken(agentData.token); + }, []); + + const onSaveClicked = useCallback(() => { + if (!token) return; + + updateAgentFromLocalStorage({ + provider: "openai", + model: "gpt-4o-mini", + token, + }); + + toast("Setting saved!"); + }, [token]); + + return ( + +
+

Local Setting

+ +

+ Bring your OpenAI token to enable the AI assistant. Your token is + stored in localStorage. We do not store your token on our server. +

+ + + +
+ +
+
+ + ); +} diff --git a/src/app/(outerbase)/nav-layout.tsx b/src/app/(outerbase)/nav-layout.tsx index f581536e..a65a0dda 100644 --- a/src/app/(outerbase)/nav-layout.tsx +++ b/src/app/(outerbase)/nav-layout.tsx @@ -26,19 +26,10 @@ export default function NavigationLayout({ children }: PropsWithChildren) { - {/*
- } - /> -
*/} -
+ { + router.push(`/local-setting`); + }} + > + Local Setting + + + {session && ( (); const [name, setName] = useState(""); - const [source, setSource] = useState(); + const [credential, setCredential] = useState(); useEffect(() => { if (!workspaceId) return; @@ -34,24 +38,32 @@ export default function OuterbaseSourcePage() { getOuterbaseBase(workspaceId, baseId).then((base) => { if (!base) return; - setSource(base.sources[0]); + setName(base.name); + getOuterbaseBaseCredential(workspaceId, base.sources[0]?.id ?? "").then( + setCredential + ); }); }, [workspaceId, baseId]); const savedDocDriver = useMemo(() => { - if (!workspaceId || !source?.id || !baseId) return null; - return new OuterbaseQueryDriver(workspaceId, baseId, source.id); - }, [workspaceId, baseId, source?.id]); + if (!workspaceId || !credential?.id || !baseId) return null; + return new OuterbaseQueryDriver(workspaceId, baseId, credential.id); + }, [workspaceId, baseId, credential?.id]); + + // We need to send analytics to update the last used time + useEffect(() => { + sendOuterbaseBaseAnalytics(workspaceId, baseId).then().catch(); + }, [workspaceId, baseId]); const [outerbaseDriver, extensions] = useMemo(() => { - if (!workspaceId || !source) return [null, null]; + if (!workspaceId || !credential) return [null, null]; - const dialect = source.type; + const dialect = credential.type; const outerbaseConfig = { workspaceId, - sourceId: source.id, - baseId: source.base_id, + sourceId: credential.id, + baseId, token: localStorage.getItem("ob-token") ?? "", }; @@ -70,7 +82,7 @@ export default function OuterbaseSourcePage() { ]; } else if (dialect === "mysql") { return [ - new OuterbaseMySQLDriver(outerbaseConfig), + new OuterbaseMySQLDriver(outerbaseConfig, credential.database), new StudioExtensionManager([ ...createMySQLExtensions(), ...outerbaseSpecifiedDrivers, @@ -85,7 +97,7 @@ export default function OuterbaseSourcePage() { ...outerbaseSpecifiedDrivers, ]), ]; - }, [workspaceId, source]); + }, [workspaceId, credential, baseId]); if (!outerbaseDriver || !savedDocDriver) { return Loading Base ...; diff --git a/src/app/(theme)/client/s/[[...driver]]/page-client.tsx b/src/app/(theme)/client/s/[[...driver]]/page-client.tsx index 2fcb2a9e..a1f6be46 100644 --- a/src/app/(theme)/client/s/[[...driver]]/page-client.tsx +++ b/src/app/(theme)/client/s/[[...driver]]/page-client.tsx @@ -3,17 +3,31 @@ import { updateLocalConnectionUsed, useLocalConnection, } from "@/app/(outerbase)/local/hooks"; -import MyStudio from "@/components/my-studio"; +import { Studio } from "@/components/gui/studio"; +import { StudioExtensionManager } from "@/core/extension-manager"; +import { + createMySQLExtensions, + createPostgreSQLExtensions, + createSQLiteExtensions, + createStandardExtensions, +} from "@/core/standard-extension"; import { createLocalDriver } from "@/drivers/helpers"; import IndexdbSavedDoc from "@/drivers/saved-doc/indexdb-saved-doc"; -import { useSearchParams } from "next/navigation"; -import { useEffect, useMemo } from "react"; +import { useAgentFromLocalStorage } from "@/lib/ai-agent-storage"; +import { useRouter, useSearchParams } from "next/navigation"; +import { useCallback, useEffect, useMemo } from "react"; export default function ClientPageBody() { const params = useSearchParams(); const baseId = params.get("p") ?? ""; const { data: conn } = useLocalConnection(baseId); + const router = useRouter(); + + const goBack = useCallback(() => { + router.push("/"); + }, [router]); + useEffect(() => { if (!baseId) return; updateLocalConnectionUsed(baseId).then().catch(); @@ -28,22 +42,42 @@ export default function ClientPageBody() { return createLocalDriver(config); }, [conn]); + const extensions = useMemo(() => { + if (!driver) return null; + const dialet = driver.getFlags().dialect; + + if (dialet === "mysql") { + return new StudioExtensionManager(createMySQLExtensions()); + } else if (dialet === "sqlite") { + return new StudioExtensionManager(createSQLiteExtensions()); + } else if (dialet === "postgres") { + return new StudioExtensionManager(createPostgreSQLExtensions()); + } + + return new StudioExtensionManager(createStandardExtensions()); + }, [driver]); + + const agentDriver = useAgentFromLocalStorage(driver); + const docDriver = useMemo(() => { if (conn) { return new IndexdbSavedDoc(conn.id); } }, [conn]); - if (!driver || !conn) { + if (!driver || !conn || !extensions) { return
Something wrong
; } return ( - ); } diff --git a/src/app/(theme)/embed/[driver]/page-client.tsx b/src/app/(theme)/embed/[driver]/page-client.tsx index bf3f533c..02244520 100644 --- a/src/app/(theme)/embed/[driver]/page-client.tsx +++ b/src/app/(theme)/embed/[driver]/page-client.tsx @@ -13,6 +13,7 @@ import { } from "@/drivers/iframe-driver"; import ElectronSavedDocs from "@/drivers/saved-doc/electron-saved-doc"; import DoltExtension from "@/extensions/dolt"; +import { useAgentFromLocalStorage } from "@/lib/ai-agent-storage"; import { useSearchParams } from "next/navigation"; import { useEffect, useMemo } from "react"; @@ -37,6 +38,8 @@ export default function EmbedPageClient({ return new StudioExtensionManager(createEmbedExtensions(driverName)); }, [driverName]); + const agentDriver = useAgentFromLocalStorage(driver); + useEffect(() => { return driver.listen(); }, [driver]); @@ -48,6 +51,7 @@ export default function EmbedPageClient({ docDriver={savedDocDriver} name={searchParams.get("name") || "Unnamed Connection"} color={searchParams.get("color") || "gray"} + agentDriver={agentDriver} /> ); } diff --git a/src/app/(theme)/playground/client/page-client.tsx b/src/app/(theme)/playground/client/page-client.tsx index 3c56547c..e91d366c 100644 --- a/src/app/(theme)/playground/client/page-client.tsx +++ b/src/app/(theme)/playground/client/page-client.tsx @@ -10,6 +10,7 @@ import { StudioExtensionManager } from "@/core/extension-manager"; import { createSQLiteExtensions } from "@/core/standard-extension"; import SqljsDriver from "@/drivers/sqljs-driver"; import { localDb } from "@/indexdb"; +import { useAgentFromLocalStorage } from "@/lib/ai-agent-storage"; import downloadFileFromUrl from "@/lib/download-file"; import { saveAs } from "file-saver"; import { @@ -43,6 +44,8 @@ export default function PlaygroundEditorBody({ const [handler, setHandler] = useState(); const [fileName, setFilename] = useState(""); + const agentDriver = useAgentFromLocalStorage(driver); + /** * Initialize the SQL.js library. */ @@ -274,12 +277,13 @@ export default function PlaygroundEditorBody({ name="Playground" driver={driver} containerClassName="w-full h-full" + agentDriver={agentDriver} /> ); } return
; - }, [databaseLoading, preloadDatabase, driver, extensions]); + }, [databaseLoading, preloadDatabase, driver, extensions, agentDriver]); return ( <> diff --git a/src/components/editor/prompt-plugin.tsx b/src/components/editor/prompt-plugin.tsx index 255d9220..a3bbc338 100644 --- a/src/components/editor/prompt-plugin.tsx +++ b/src/components/editor/prompt-plugin.tsx @@ -37,11 +37,12 @@ export interface PromptSelectedFragment { fullText: string; startLineNumber: number; endLineNumber: number; + sessionId: string; } export type PromptCallback = ( promptQuery: string, - selected?: PromptSelectedFragment + selected: PromptSelectedFragment ) => Promise; class PlaceholderWidget extends WidgetType { @@ -69,6 +70,9 @@ class PromptWidget extends WidgetType { ) { super(); + // Generate unique session id for this prompt + const sessionId = crypto.randomUUID(); + plugin.lock(); this.container = document.createElement("div"); @@ -84,6 +88,7 @@ class PromptWidget extends WidgetType { const getSelectionLines = () => { const startLineNumber = view.state.doc.lineAt(from).number; const endLineNumber = view.state.doc.lineAt(to).number; + return Array.from( { length: endLineNumber - startLineNumber + 1 }, (_, i) => startLineNumber + i @@ -150,6 +155,7 @@ class PromptWidget extends WidgetType { text: suggestedText ?? selectedOriginalText, fullText: view.state.doc.toString(), startLineNumber, + sessionId, endLineNumber: view.state.doc.lineAt( startPosition + suggestedText.length ).number, @@ -555,7 +561,7 @@ export class CodeMirrorPromptPlugin { ]; } - async getSuggestion(promptQuery: string, selected?: PromptSelectedFragment) { + async getSuggestion(promptQuery: string, selected: PromptSelectedFragment) { if (this.promptCallback) { return this.promptCallback(promptQuery, selected); } diff --git a/src/components/gui/studio.tsx b/src/components/gui/studio.tsx index d2db2998..b181395a 100644 --- a/src/components/gui/studio.tsx +++ b/src/components/gui/studio.tsx @@ -4,6 +4,7 @@ import { ConfigProvider } from "@/context/config-provider"; import { DriverProvider } from "@/context/driver-provider"; import { StudioExtensionManager } from "@/core/extension-manager"; import { BeforeQueryPipeline } from "@/core/query-pipeline"; +import { AgentBaseDriver } from "@/drivers/agent/base"; import type { BaseDriver } from "@/drivers/base-driver"; import { CollaborationBaseDriver } from "@/drivers/collaboration-driver-base"; import { SavedDocDriver } from "@/drivers/saved-doc/saved-doc-driver"; @@ -15,6 +16,7 @@ interface StudioProps { driver: BaseDriver; extensions?: StudioExtensionManager; collaboration?: CollaborationBaseDriver; + agentDriver?: AgentBaseDriver; docDriver?: SavedDocDriver; name: string; color: string; @@ -32,6 +34,7 @@ export function Studio({ extensions, containerClassName, docDriver, + agentDriver, }: Readonly) { const extensionRef = useRef( extensions @@ -92,8 +95,16 @@ export function Studio({ onBack, extensions: finalExtensionManager, containerClassName, + agentDriver, }; - }, [name, color, onBack, finalExtensionManager, containerClassName]); + }, [ + name, + color, + onBack, + finalExtensionManager, + containerClassName, + agentDriver, + ]); return ( (initialSavedKey); const [placeholders, setPlaceholders] = useState>({}); + const { agentDriver } = useConfig(); + const { schema } = useSchema(); + useEffect(() => { const timer = setTimeout(() => { setPlaceholders((prev) => { @@ -317,6 +322,26 @@ export default function QueryWindow({ [] ); + const promptConversationIds = useRef>({}); + const onPrompt = useCallback( + async (promptQuery: string, option: PromptSelectedFragment) => { + if (!agentDriver) return ""; + + const agentResponse = await agentDriver.promptInline( + promptQuery, + promptConversationIds.current[option.sessionId], + { + selected: option?.text ?? "", + schema: schema, + } + ); + + promptConversationIds.current[option.sessionId] = agentResponse.id; + return agentResponse.result; + }, + [agentDriver, schema] + ); + return ( @@ -394,6 +419,8 @@ export default function QueryWindow({
{ - router.push("/"); - }, [router]); - - const extensions = useMemo(() => { - if (dialet === "mysql") { - return new StudioExtensionManager(createMySQLExtensions()); - } else if (dialet === "sqlite") { - return new StudioExtensionManager(createSQLiteExtensions()); - } else if (dialet === "postgres") { - return new StudioExtensionManager(createPostgreSQLExtensions()); - } - - return new StudioExtensionManager(createStandardExtensions()); - }, [dialet]); - - return ( - - ); -} - -export default function MyStudio(props: MyStudioProps) { - return ; -} diff --git a/src/context/config-provider.tsx b/src/context/config-provider.tsx index b1f78583..575dc4dd 100644 --- a/src/context/config-provider.tsx +++ b/src/context/config-provider.tsx @@ -1,4 +1,5 @@ import { StudioExtensionManager } from "@/core/extension-manager"; +import { AgentBaseDriver } from "@/drivers/agent/base"; import { noop } from "lodash"; import type { PropsWithChildren } from "react"; import { createContext, useContext } from "react"; @@ -9,6 +10,7 @@ interface ConfigContextProps { onBack?: () => void; extensions: StudioExtensionManager; containerClassName?: string; + agentDriver?: AgentBaseDriver; } const ConfigContext = createContext({ diff --git a/src/context/connection-config-provider.tsx b/src/context/connection-config-provider.tsx deleted file mode 100644 index 008e1d24..00000000 --- a/src/context/connection-config-provider.tsx +++ /dev/null @@ -1,25 +0,0 @@ -import { SavedConnectionItem } from "@/app/(theme)/connect/saved-connection-storage"; -import { createContext, PropsWithChildren, useContext, useMemo } from "react"; - -const ConnectionConfigContext = createContext<{ - config?: SavedConnectionItem; -}>({}); - -export function useConnectionConfig() { - const { config } = useContext(ConnectionConfigContext); - if (!config) throw new Error("Need Connection Config Provider"); - return { config }; -} - -export function ConnectionConfigProvider({ - children, - config, -}: PropsWithChildren<{ config: SavedConnectionItem }>) { - const contextMemo = useMemo(() => ({ config }), [config]); - - return ( - - {children} - - ); -} diff --git a/src/drivers/agent/base.ts b/src/drivers/agent/base.ts new file mode 100644 index 00000000..379d3b89 --- /dev/null +++ b/src/drivers/agent/base.ts @@ -0,0 +1,26 @@ +import { DatabaseSchemas } from "../base-driver"; + +export interface AgentPromptOption { + schema?: DatabaseSchemas; + selectedSchema?: string; + selected: string; +} + +export interface AgentPromptResponse { + result: string; + id: string; +} + +export abstract class AgentBaseDriver { + /** + * + * @param message User message + * @param previousId Previous message id. If not provided, it is a new conversation + * @param option + */ + abstract promptInline( + message: string, + previousId: string | undefined, + option: AgentPromptOption + ): Promise; +} diff --git a/src/drivers/agent/chatgpt.ts b/src/drivers/agent/chatgpt.ts new file mode 100644 index 00000000..af9c8118 --- /dev/null +++ b/src/drivers/agent/chatgpt.ts @@ -0,0 +1,185 @@ +import { + BaseDriver, + DatabaseSchemas, + DatabaseTableSchema, +} from "../base-driver"; +import { + AgentBaseDriver, + AgentPromptOption, + AgentPromptResponse, +} from "./base"; + +interface ChatHistory { + id: string; + createdAt: number; + messages: { role: string; content: string }[]; +} + +interface ChatGPTResponse { + choices: { message: { role: string; content: string } }[]; +} + +export class ChatGPTDriver implements AgentBaseDriver { + protected history: Record = {}; + + constructor( + protected driver: BaseDriver, + protected token: string + ) {} + + protected convertTableToContent( + schemaName: string | undefined, + table: DatabaseTableSchema + ): string { + const columns = table.columns + .map((column) => { + return `${this.driver.escapeId(column.name)} ${column.type}`; + }) + .join(",\n"); + + const fullTableName = schemaName + ? `${this.driver.escapeId(schemaName)}.${this.driver.escapeId(table.tableName ?? "")}` + : this.driver.escapeId(table.tableName ?? ""); + + const primaryKeyPart = + table.pk.length > 0 + ? `, PRIMARY KEY (${table.pk.map(this.driver.escapeId).join(", ")})` + : ""; + + const foreignKeyPart: string[] = []; + for (const column of table.columns) { + if (column.constraint?.foreignKey) { + foreignKeyPart.push( + [ + "FOREIGN KEY", + column.name, + "REFERENCES", + column.constraint.foreignKey.foreignTableName ?? "", + "(", + (column.constraint?.foreignKey?.foreignColumns ?? [])[0] ?? "", + ")", + ].join(" ") + ); + } + } + + for (const constraint of table.constraints ?? []) { + if (constraint.foreignKey) { + foreignKeyPart.push( + [ + "FOREIGN KEY", + `(${(constraint.foreignKey.columns ?? []).join(", ")})`, + "REFERENCES", + constraint.foreignKey.foreignTableName ?? "", + `(${(constraint.foreignKey.foreignColumns ?? []).join(", ")})`, + ].join(" ") + ); + } + } + + return `CREATE TABLE ${fullTableName} (\n${columns}\n ${primaryKeyPart});`; + } + + protected convertSchemaToContent(schemas: DatabaseSchemas): string { + const schemaParts: string[] = []; + const defaultSchema = this.driver.getFlags().defaultSchema; + + for (const [schemaName, schema] of Object.entries(schemas)) { + for (const table of schema) { + if (!table.tableSchema) continue; + if (!["table", "view"].includes(table.type)) continue; + + schemaParts.push( + this.convertTableToContent( + defaultSchema.toLowerCase() === schemaName.toLowerCase() + ? "" + : schemaName, + table.tableSchema + ) + ); + } + } + + return schemaParts.join("\n\n"); + } + + async promptInline( + message: string, + previousId: string | undefined, + option: AgentPromptOption + ): Promise { + const session = this.history[previousId ?? ""] ?? { + id: crypto.randomUUID(), + createdAt: Date.now(), + messages: [], + }; + + if (session.messages.length === 0) { + session.messages.push({ + role: "system", + content: "You are an SQL expert. Only return SQL code.", + }); + + session.messages.push({ + role: "user", + content: + "Here is " + + this.driver.getFlags().dialect + + " my database schema:\n\n", + }); + + if (option.schema) { + session.messages.push({ + role: "user", + content: + "```sql\n" + this.convertSchemaToContent(option.schema) + "```", + }); + } + + if (option.selected) { + session.messages.push({ + role: "user", + content: + "This is my selected query ```sql\n" + option.selected + "```", + }); + } + } + + session.messages.push({ + role: "user", + content: message, + }); + + const response = await fetch("https://api.openai.com/v1/chat/completions", { + method: "POST", + headers: { + Authorization: `Bearer ${this.token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + model: "gpt-4o-mini-2024-07-18", + temperature: 0, + messages: session.messages, + }), + }); + + const jsonResponse = (await response.json()) as ChatGPTResponse; + const suggestedQuery = jsonResponse.choices[0].message.content; + + // Striped the SQL code from the response + const sqlCode = suggestedQuery.replace(/```sql\n/g, "").replace(/```/g, ""); + + // Save the chat history + session.messages.push({ + role: "assistant", + content: suggestedQuery, + }); + + this.history[session.id] = session; + + return { + id: session.id, + result: sqlCode, + }; + } +} diff --git a/src/drivers/mysql/mysql-driver.ts b/src/drivers/mysql/mysql-driver.ts index 2f133d7d..dcfd18bb 100644 --- a/src/drivers/mysql/mysql-driver.ts +++ b/src/drivers/mysql/mysql-driver.ts @@ -128,6 +128,11 @@ function mapColumn(column: MySqlColumn): DatabaseTableColumn { export default abstract class MySQLLikeDriver extends CommonSQLImplement { columnTypeSelector: ColumnTypeSelector = MYSQL_DATA_TYPE_SUGGESTION; + // If this is specified, we only show the tables in this database + // Outerbase Cloud does not support the USE statement because it runs in non-interactive mode + // It does not make sense to show other databases. + selectedDatabase: string = ""; + escapeId(id: string) { return `\`${id.replace(/`/g, "``")}\``; } @@ -138,14 +143,13 @@ export default abstract class MySQLLikeDriver extends CommonSQLImplement { getFlags(): DriverFlags { return { - defaultSchema: "", - optionalSchema: false, + defaultSchema: this.selectedDatabase, + optionalSchema: this.selectedDatabase ? true : false, supportBigInt: false, supportModifyColumn: true, supportCreateUpdateTable: true, - supportCreateUpdateDatabase: true, + supportCreateUpdateDatabase: this.selectedDatabase ? false : true, dialect: "mysql", - supportUseStatement: true, supportRowId: false, supportInsertReturning: false, @@ -167,22 +171,29 @@ export default abstract class MySQLLikeDriver extends CommonSQLImplement { } async schemas(): Promise { - const schemaSql = - "SELECT SCHEMA_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; + const schemaSql = this.selectedDatabase + ? `SELECT SCHEMA_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = ${this.escapeValue(this.selectedDatabase)}` + : "SELECT SCHEMA_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; - const tableSql = - "SELECT TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE, DATA_LENGTH, INDEX_LENGTH FROM information_schema.tables WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; + const tableSql = this.selectedDatabase + ? `SELECT TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE, DATA_LENGTH, INDEX_LENGTH FROM information_schema.tables WHERE TABLE_SCHEMA = ${this.escapeValue(this.selectedDatabase)}` + : "SELECT TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE, DATA_LENGTH, INDEX_LENGTH FROM information_schema.tables WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; - const columnSql = - "SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, COLUMN_TYPE, DATA_TYPE, EXTRA, COLUMN_KEY, IS_NULLABLE, COLUMN_DEFAULT FROM information_schema.columns WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; + const columnSql = this.selectedDatabase + ? `SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, COLUMN_TYPE, DATA_TYPE, EXTRA, COLUMN_KEY, IS_NULLABLE, COLUMN_DEFAULT FROM information_schema.columns WHERE TABLE_SCHEMA = ${this.escapeValue(this.selectedDatabase)}` + : "SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, COLUMN_TYPE, DATA_TYPE, EXTRA, COLUMN_KEY, IS_NULLABLE, COLUMN_DEFAULT FROM information_schema.columns WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; - const constraintSql = - "SELECT TABLE_SCHEMA, TABLE_NAME, CONSTRAINT_NAME, CONSTRAINT_TYPE FROM information_schema.table_constraints WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys') AND CONSTRAINT_TYPE IN ('PRIMARY KEY', 'UNIQUE', 'FOREIGN KEY')"; + const constraintSql = this.selectedDatabase + ? `SELECT TABLE_SCHEMA, TABLE_NAME, CONSTRAINT_NAME, CONSTRAINT_TYPE FROM information_schema.table_constraints WHERE TABLE_SCHEMA = ${this.escapeValue(this.selectedDatabase)} AND CONSTRAINT_TYPE IN ('PRIMARY KEY', 'UNIQUE', 'FOREIGN KEY')` + : "SELECT TABLE_SCHEMA, TABLE_NAME, CONSTRAINT_NAME, CONSTRAINT_TYPE FROM information_schema.table_constraints WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys') AND CONSTRAINT_TYPE IN ('PRIMARY KEY', 'UNIQUE', 'FOREIGN KEY')"; - const constraintColumnsSql = `SELECT CONSTRAINT_NAME, TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME FROM information_schema.key_column_usage WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')`; + const constraintColumnsSql = this.selectedDatabase + ? `SELECT CONSTRAINT_NAME, TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME FROM information_schema.key_column_usage WHERE TABLE_SCHEMA = ${this.escapeValue(this.selectedDatabase)}` + : `SELECT CONSTRAINT_NAME, TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME, REFERENCED_TABLE_SCHEMA, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME FROM information_schema.key_column_usage WHERE TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')`; - const triggerSql = - "SELECT * from information_schema.triggers WHERE TRIGGER_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; + const triggerSql = this.selectedDatabase + ? `SELECT * from information_schema.triggers WHERE TRIGGER_SCHEMA = ${this.escapeValue(this.selectedDatabase)}` + : "SELECT * from information_schema.triggers WHERE TRIGGER_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"; const result = await this.batch([ schemaSql, diff --git a/src/hooks/use-connect.ts b/src/hooks/use-connect.ts deleted file mode 100644 index 6fa8f18d..00000000 --- a/src/hooks/use-connect.ts +++ /dev/null @@ -1,30 +0,0 @@ -"use client"; -import { - SavedConnectionItemConfigConfig, - SupportedDriver, -} from "@/app/(theme)/connect/saved-connection-storage"; -import { useRouter } from "next/navigation"; -import { useCallback } from "react"; - -export default function useConnect() { - const router = useRouter(); - - return useCallback( - (driver: SupportedDriver, config: SavedConnectionItemConfigConfig) => { - sessionStorage.setItem( - "connection", - JSON.stringify({ - driver, - url: config.url, - token: config.token, - username: config.username, - password: config.password, - database: config.database, - }) - ); - - router.push(`/client/${driver ?? "turso"}`); - }, - [router] - ); -} diff --git a/src/lib/ai-agent-storage.ts b/src/lib/ai-agent-storage.ts new file mode 100644 index 00000000..28fd34d0 --- /dev/null +++ b/src/lib/ai-agent-storage.ts @@ -0,0 +1,42 @@ +import { ChatGPTDriver } from "@/drivers/agent/chatgpt"; +import { BaseDriver } from "@/drivers/base-driver"; +import { useMemo } from "react"; + +export interface LocalAgentType { + provider: "openai"; + model: "gpt-4o-mini"; + token: string; +} + +export function getAgentFromLocalStorage(): LocalAgentType | undefined { + if (typeof window === "undefined") return undefined; + + // Getting the driver from the local storage + const agentRawData = localStorage.getItem("agent"); + if (!agentRawData) return undefined; + + // Parsing the agent data + const agentData: LocalAgentType = JSON.parse(agentRawData); + + // Validate the data + if (agentData.provider !== "openai") return undefined; + if (agentData.model !== "gpt-4o-mini") return undefined; + if (!agentData.token) return undefined; + + return agentData; +} + +export function updateAgentFromLocalStorage(data: LocalAgentType) { + localStorage.setItem("agent", JSON.stringify(data)); +} + +export function useAgentFromLocalStorage(databaseDriver?: BaseDriver | null) { + return useMemo(() => { + if (!databaseDriver) return undefined; + + const agentConfig = getAgentFromLocalStorage(); + if (!agentConfig) return undefined; + + return new ChatGPTDriver(databaseDriver, agentConfig.token); + }, [databaseDriver]); +} diff --git a/src/outerbase-cloud/api-type.ts b/src/outerbase-cloud/api-type.ts index 441a6ac0..85d8c853 100644 --- a/src/outerbase-cloud/api-type.ts +++ b/src/outerbase-cloud/api-type.ts @@ -62,6 +62,10 @@ export interface OuterbaseAPISourceInput { connection_id?: string; } +export interface OuterbaseAPIBaseCredential extends OuterbaseAPISourceInput { + id: string; +} + export interface OuterbaseAPISource { model: "source"; type: string; diff --git a/src/outerbase-cloud/api-workspace.ts b/src/outerbase-cloud/api-workspace.ts index 1bab7077..9b5a1f9a 100644 --- a/src/outerbase-cloud/api-workspace.ts +++ b/src/outerbase-cloud/api-workspace.ts @@ -2,6 +2,7 @@ import { mutate } from "swr"; import { requestOuterbase } from "./api"; import { OuterbaseAPIBase, + OuterbaseAPIBaseCredential, OuterbaseAPIConnection, OuterbaseAPISource, OuterbaseAPISourceInput, @@ -103,7 +104,7 @@ export async function getOuterbaseBaseCredential( workspaceId: string, sourceId: string ) { - return await requestOuterbase( + return await requestOuterbase( `/api/v1/workspace/${workspaceId}/source/${sourceId}/credential`, "GET" ); diff --git a/src/outerbase-cloud/api.ts b/src/outerbase-cloud/api.ts index 6c0dea4f..383a7947 100644 --- a/src/outerbase-cloud/api.ts +++ b/src/outerbase-cloud/api.ts @@ -263,3 +263,19 @@ export async function getOuterbaseEmbedChart( return await result.json(); } + +export async function sendOuterbaseBaseAnalytics( + workspaceId: string, + baseId: string +) { + return requestOuterbase( + `/api/v1/workspace/${workspaceId}/base/${baseId}/analytics`, + "POST", + { + data: { + path: "/[workspaceId]/[baseId]/settings/database", + }, + type: "page_view", + } + ); +} diff --git a/src/outerbase-cloud/database/mysql.ts b/src/outerbase-cloud/database/mysql.ts index 88ba21cd..d675bda4 100644 --- a/src/outerbase-cloud/database/mysql.ts +++ b/src/outerbase-cloud/database/mysql.ts @@ -15,9 +15,13 @@ export class OuterbaseMySQLDriver extends MySQLLikeDriver { }; } - constructor({ workspaceId, sourceId }: OuterbaseDatabaseConfig) { + constructor( + { workspaceId, sourceId }: OuterbaseDatabaseConfig, + selectedDatabase?: string + ) { super(); + this.selectedDatabase = selectedDatabase ?? ""; this.workspaceId = workspaceId; this.sourceId = sourceId; }