Merge pull request #381 from outerbase/trigger-ob-schema-update

update outerbase schema when schema changed

src/app/(outerbase)/w/[workspaceId]/[baseId]/page.tsx

@@ -9,6 +9,7 @@ import {
   createSQLiteExtensions,
 } from "@/core/standard-extension";
 import DataCatalogExtension from "@/extensions/data-catalog";
+import OuterbaseExtension from "@/extensions/outerbase";
 import { getOuterbaseBase } from "@/outerbase-cloud/api";
 import { OuterbaseAPISource } from "@/outerbase-cloud/api-type";
 import DataCatalogOuterbaseDriver from "@/outerbase-cloud/data-catalog-driver";
@@ -50,11 +51,12 @@ export default function OuterbaseSourcePage() {
     const outerbaseConfig = {
       workspaceId,
       sourceId: source.id,
-      baseId,
+      baseId: source.base_id,
       token: localStorage.getItem("ob-token") ?? "",
     };
 
     const outerbaseSpecifiedDrivers = [
+      new OuterbaseExtension(outerbaseConfig),
       new DataCatalogExtension(new DataCatalogOuterbaseDriver(outerbaseConfig)),
     ];
 
@@ -83,7 +85,7 @@ export default function OuterbaseSourcePage() {
         ...outerbaseSpecifiedDrivers,
       ]),
     ];
-  }, [workspaceId, source, baseId]);
+  }, [workspaceId, source]);
 
   if (!outerbaseDriver || !savedDocDriver) {
     return <PageLoading>Loading Base ...</PageLoading>;

src/components/gui/schema-editor/schema-save-dialog.tsx

@@ -5,20 +5,20 @@ import {
   AlertDialogFooter,
   AlertDialogTitle,
 } from "@/components/ui/alert-dialog";
-import CodePreview from "../code-preview";
 import { Button } from "@/components/ui/button";
+import { useDatabaseDriver } from "@/context/driver-provider";
+import { useSchema } from "@/context/schema-provider";
+import { DatabaseTableSchemaChange } from "@/drivers/base-driver";
 import {
   LucideAlertCircle,
   LucideLoader,
   LucideSave,
   LucideTableProperties,
 } from "lucide-react";
-import { useTabsContext } from "../windows-tab";
-import { useDatabaseDriver } from "@/context/driver-provider";
-import { useSchema } from "@/context/schema-provider";
 import { useCallback, useState } from "react";
+import CodePreview from "../code-preview";
 import SchemaEditorTab from "../tabs/schema-editor-tab";
-import { DatabaseTableSchemaChange } from "@/drivers/base-driver";
+import { useTabsContext } from "../windows-tab";
 
 export default function SchemaSaveDialog({
   schema,
@@ -42,8 +42,9 @@ export default function SchemaSaveDialog({
     databaseDriver
       .transaction(previewScript)
       .then(() => {
+        refreshSchema();
+
         if (schema.name.new !== schema.name.old) {
-          refreshSchema();
           replaceCurrentTab({
             component: (
               <SchemaEditorTab
@@ -83,8 +84,8 @@ export default function SchemaSaveDialog({
         <AlertDialogTitle>Preview</AlertDialogTitle>
 
         {errorMessage && (
-          <div className="text-sm text-red-500 font-mono flex gap-4 justify-end items-end">
-            <LucideAlertCircle className="w-12 h-12" />
+          <div className="flex items-end justify-end gap-4 font-mono text-sm text-red-500">
+            <LucideAlertCircle className="h-12 w-12" />
             <p>{errorMessage}</p>
           </div>
         )}
@@ -95,9 +96,9 @@ export default function SchemaSaveDialog({
           <AlertDialogCancel>Cancel</AlertDialogCancel>
           <Button onClick={onSave}>
             {isExecuting ? (
-              <LucideLoader className="w-4 h-4 mr-2 animate-spin" />
+              <LucideLoader className="mr-2 h-4 w-4 animate-spin" />
             ) : (
-              <LucideSave className="w-4 h-4 mr-2" />
+              <LucideSave className="mr-2 h-4 w-4" />
             )}
             Continue
           </Button>

src/context/schema-provider.tsx

@@ -10,6 +10,7 @@ import {
   useState,
 } from "react";
 import { useAutoComplete } from "./auto-complete-provider";
+import { useConfig } from "./config-provider";
 import { useDatabaseDriver } from "./driver-provider";
 
 type AutoCompletionSchema = Record<string, Record<string, string[]> | string[]>;
@@ -68,6 +69,7 @@ export function SchemaProvider({ children }: Readonly<PropsWithChildren>) {
   const [error, setError] = useState<string>();
   const [loading, setLoading] = useState(true);
   const { databaseDriver } = useDatabaseDriver();
+  const { extensions } = useConfig();
 
   const [schema, setSchema] = useState<DatabaseSchemas>({});
   const [currentSchema, setCurrentSchema] = useState<DatabaseSchemaItem[]>([]);
@@ -119,6 +121,18 @@ export function SchemaProvider({ children }: Readonly<PropsWithChildren>) {
     }
   }, [currentSchemaName, schema, setCurrentSchema]);
 
+  /**
+   * Triggered when re-fetching the database schema.
+   * This is particularly useful for Outerbase Cloud,
+   * which needs to update its data catalog to provide
+   * the schema to the AI.
+   */
+  useEffect(() => {
+    if (schema && Object.entries(schema).length > 0) {
+      extensions.triggerAfterFetchSchemaCallback(schema);
+    }
+  }, [schema, extensions]);
+
   useEffect(() => {
     const sortedTableList = [...currentSchema];
     sortedTableList.sort((a, b) => {

src/core/extension-manager.tsx

@@ -1,5 +1,5 @@
 import { OptimizeTableHeaderProps } from "@/components/gui/table-optimized";
-import { DatabaseSchemaItem } from "@/drivers/base-driver";
+import { DatabaseSchemaItem, DatabaseSchemas } from "@/drivers/base-driver";
 import { ReactElement } from "react";
 import { IStudioExtension } from "./extension-base";
 import { BeforeQueryPipeline } from "./query-pipeline";
@@ -31,13 +31,16 @@ type QueryHeaderResultMenuHandler = (
   header: OptimizeTableHeaderProps
 ) => StudioExtensionMenuItem | undefined;
 
+type AfterFetchSchemaHandler = (schema: DatabaseSchemas) => void;
+
 type QueryResultCellMenuHandler = () => StudioExtensionMenuItem | undefined;
 
 export class StudioExtensionContext {
   protected sidebars: RegisterSidebarOption[] = [];
 
   protected beforeQueryHandlers: BeforeQueryHandler[] = [];
   protected afterQueryHandlers: AfterQueryHandler[] = [];
+  protected afterFetchSchemaHandlers: AfterFetchSchemaHandler[] = [];
 
   protected queryResultHeaderContextMenu: QueryHeaderResultMenuHandler[] = [];
   protected queryResultCellContextMenu: QueryResultCellMenuHandler[] = [];
@@ -58,6 +61,10 @@ export class StudioExtensionContext {
     this.afterQueryHandlers.push(handler);
   }
 
+  registerAfterFetchSchema(handler: AfterFetchSchemaHandler) {
+    this.afterFetchSchemaHandlers.push(handler);
+  }
+
   registerSidebar(option: RegisterSidebarOption) {
     this.sidebars.push(option);
   }
@@ -125,6 +132,12 @@ export class StudioExtensionManager extends StudioExtensionContext {
       .filter(Boolean) as StudioExtensionMenuItem[];
   }
 
+  triggerAfterFetchSchemaCallback(schema: DatabaseSchemas) {
+    for (const handler of this.afterFetchSchemaHandlers) {
+      handler(schema);
+    }
+  }
+
   async beforeQuery(payload: BeforeQueryPipeline) {
     for (const handler of this.beforeQueryHandlers) {
       await handler(payload);

src/extensions/data-catalog/table-result-header.tsx

@@ -22,17 +22,16 @@ export default function DataCatalogResultHeader({
   const [loading, setLoading] = useState(false);
 
   const onSaveClicked = useCallback(() => {
-    if (!column) return;
     setLoading(true);
 
     driver
       .updateColumn(schemaName, tableName, columnName, {
         definition,
-        hide: column.hide || false,
-        samples: column.samples,
+        hide: column?.hide || false,
+        samples: column?.samples ?? [],
       })
       .then(() => {
-        toast.success("updated");
+        toast.success("Column Definition Updated");
       })
       .catch()
       .finally(() => setLoading(false));

src/extensions/outerbase/index.tsx

@@ -0,0 +1,35 @@
+/**
+ * Outerbase Cloud extension for studio functionality.
+ */
+import { StudioExtension } from "@/core/extension-base";
+import { StudioExtensionContext } from "@/core/extension-manager";
+import { OuterbaseDatabaseConfig } from "@/outerbase-cloud/api-type";
+import { updateOuterbaseSchemas } from "@/outerbase-cloud/api-workspace";
+
+export default class OuterbaseExtension extends StudioExtension {
+  extensionName = "outerbase-cloud-studio";
+
+  constructor(protected config: OuterbaseDatabaseConfig) {
+    super();
+  }
+
+  init(studio: StudioExtensionContext): void {
+    // The Outerbase Data Catalog and AI-powered features
+    // require the schema to be stored in the Outerbase API database.
+    // Originally, Outerbase fetched the schema from an API endpoint
+    // and stored it in its database.
+    //
+    // Our studio differs because we obtain the schema from
+    // raw queries. The Outerbase API does not know when to refetch
+    // the updated schema and store it in its database.
+    // We hook into the event when the studio finishes fetching the schema
+    // and notify Outerbase Cloud to refetch the schema on their side.
+    //
+    // TECHNICAL DEBT: This is redundant work and should be optimized in the future.
+    studio.registerAfterFetchSchema(() => {
+      updateOuterbaseSchemas(this.config.workspaceId, this.config.sourceId)
+        .then()
+        .catch();
+    });
+  }
+}

src/outerbase-cloud/api-data-catalog.ts

@@ -11,11 +11,11 @@ import {
 
 export async function getOuterbaseSchemas(
   workspaceId: string,
-  sourceId: string,
-  baseId?: string
+  sourceId: string
 ) {
+  // 
   return await requestOuterbase<Record<string, OuterbaseDataCatalogSchemas[]>>(
-    `/api/v1/workspace/${workspaceId}/source/${sourceId}/schema?baseId=${baseId}`
+    `/api/v1/workspace/${workspaceId}/source/${sourceId}/schema?baseId=`
   );
 }
 

src/outerbase-cloud/api-type.ts

@@ -66,6 +66,7 @@ export interface OuterbaseAPISource {
   model: "source";
   type: string;
   id: string;
+  base_id: string;
 }
 export interface OuterbaseAPIBase {
   model: "base";
@@ -151,7 +152,7 @@ export interface OuterbaseAPIUser {
 
 export interface OuterbaseAPIDashboardDetail
   extends OuterbaseAPIDashboard,
-    DashboardProps {}
+  DashboardProps { }
 
 export interface OuterbaseAPIWorkspaceResponse {
   items: OuterbaseAPIWorkspace[];

src/outerbase-cloud/api-workspace.ts

@@ -75,3 +75,13 @@ export function createOuterbaseSource(
     source
   );
 }
+
+export async function updateOuterbaseSchemas(
+  workspaceId: string,
+  sourceId: string
+) {
+  // 
+  return await requestOuterbase<unknown>(
+    `/api/v1/workspace/${workspaceId}/source/${sourceId}/schema?baseId=`, "POST"
+  );
+}