From 53b782dc0f60a032c66b74c81543f71e831dd88b Mon Sep 17 00:00:00 2001
From: Thomas Draier <thomas@dust.tt>
Date: Sun, 11 Aug 2024 21:22:35 +0200
Subject: [PATCH] minor fixes

---
 front/lib/api/vaults.ts                       |  6 +-
 front/lib/swr.ts                              | 60 ++++++++++++++-----
 .../data_source_views/[dsvId]/content.ts      |  8 ++-
 types/src/front/api_handlers/public/vaults.ts | 14 ++++-
 4 files changed, 66 insertions(+), 22 deletions(-)

diff --git a/front/lib/api/vaults.ts b/front/lib/api/vaults.ts
index 2389321356ff..451801f65ca5 100644
--- a/front/lib/api/vaults.ts
+++ b/front/lib/api/vaults.ts
@@ -30,7 +30,7 @@ export type LightContentNode = {
   lastUpdatedAt: number | null;
 };
 
-export type GetDataSourceContentResponseBody = {
+export type GetDataSourceOrViewContentResponseBody = {
   nodes: LightContentNode[];
 };
 
@@ -92,7 +92,9 @@ export const getDataSourceCategory = (
 
 export const getContentHandler = async (
   req: NextApiRequest,
-  res: NextApiResponse<WithAPIErrorResponse<GetDataSourceContentResponseBody>>,
+  res: NextApiResponse<
+    WithAPIErrorResponse<GetDataSourceOrViewContentResponseBody>
+  >,
   dataSource: DataSourceResource,
   rootIds: string[] | null
 ): Promise<void> => {
diff --git a/front/lib/swr.ts b/front/lib/swr.ts
index 3e4c8afed057..fdaa05e515c9 100644
--- a/front/lib/swr.ts
+++ b/front/lib/swr.ts
@@ -18,6 +18,7 @@ import useSWR from "swr";
 import useSWRInfinite from "swr/infinite";
 
 import type { FetchConversationMessagesResponse } from "@app/lib/api/assistant/messages";
+import type { GetDataSourceOrViewContentResponseBody } from "@app/lib/api/vaults";
 import { COMMIT_HASH } from "@app/lib/commit-hash";
 import type { GetPokePlansResponseBody } from "@app/pages/api/poke/plans";
 import type { GetPokeWorkspacesResponseBody } from "@app/pages/api/poke/workspaces";
@@ -1328,18 +1329,49 @@ export function useVaultDataSourceOrViews({
     vaultsDataSourcesFetcher
   );
 
-  console.log(data, error);
-  if (data && data.dataSources) {
-    return {
-      vaultDataSources: data ? data.dataSources : null,
-      isVaultDataSourcesLoading: !error && !data,
-      isVaultDataSourcesError: error,
-    };
-  } else {
-    return {
-      vaultDataSources: data ? data.dataSourceViews : null,
-      isVaultDataSourcesLoading: !error && !data,
-      isVaultDataSourcesError: error,
-    };
-  }
+  const vaultDataSourceOrViews =
+    type === "data_sources"
+      ? (data as GetVaultDataSourcesResponseBody)?.dataSources
+      : (data as GetVaultDataSourceViewsResponseBody)?.dataSourceViews;
+
+  return {
+    vaultDataSourceOrViews,
+    isVaultDataSourceOrViewsLoading: !error && !data,
+    isVaultDataSourceOrViewsError: error,
+  };
+}
+
+export function useVaultDataSourceOrViewContent({
+  workspaceId,
+  vaultId,
+  type,
+  dataSourceOrViewId,
+  viewType,
+  parentId,
+  disabled,
+}: {
+  workspaceId: string;
+  vaultId: string;
+  type: "data_sources" | "data_source_views";
+  dataSourceOrViewId: string;
+  viewType: ContentNodesViewType;
+  parentId: string | undefined;
+  disabled?: boolean;
+}) {
+  const vaultsDataSourcesFetcher: Fetcher<GetDataSourceOrViewContentResponseBody> =
+    fetcher;
+  const qs =
+    `?viewType=${viewType}` + (parentId ? `&parentId=${parentId}` : "");
+  const { data, error } = useSWRWithDefaults(
+    disabled
+      ? null
+      : `/api/w/${workspaceId}/vaults/${vaultId}/${type}/${dataSourceOrViewId}/content${qs}`,
+    vaultsDataSourcesFetcher
+  );
+
+  return {
+    vaultContent: data?.nodes,
+    isVaultContentLoading: !error && !data,
+    isVaultContentError: error,
+  };
 }
diff --git a/front/pages/api/w/[wId]/vaults/[vId]/data_source_views/[dsvId]/content.ts b/front/pages/api/w/[wId]/vaults/[vId]/data_source_views/[dsvId]/content.ts
index aca33bc90d16..8be61bf09e58 100644
--- a/front/pages/api/w/[wId]/vaults/[vId]/data_source_views/[dsvId]/content.ts
+++ b/front/pages/api/w/[wId]/vaults/[vId]/data_source_views/[dsvId]/content.ts
@@ -1,7 +1,7 @@
 import type { WithAPIErrorResponse } from "@dust-tt/types";
 import type { NextApiRequest, NextApiResponse } from "next";
 
-import type { GetDataSourceContentResponseBody } from "@app/lib/api/vaults";
+import type { GetDataSourceOrViewContentResponseBody } from "@app/lib/api/vaults";
 import { getContentHandler } from "@app/lib/api/vaults";
 import { withSessionAuthenticationForWorkspace } from "@app/lib/api/wrappers";
 import type { Authenticator } from "@app/lib/auth";
@@ -10,7 +10,9 @@ import { apiError } from "@app/logger/withlogging";
 
 async function handler(
   req: NextApiRequest,
-  res: NextApiResponse<WithAPIErrorResponse<GetDataSourceContentResponseBody>>,
+  res: NextApiResponse<
+    WithAPIErrorResponse<GetDataSourceOrViewContentResponseBody>
+  >,
   auth: Authenticator
 ): Promise<void> {
   const owner = auth.workspace();
@@ -30,7 +32,7 @@ async function handler(
   );
 
   const dataSource = dataSourceView?.dataSource;
-  const vault = dataSource?.vault;
+  const vault = dataSourceView?.vault;
 
   if (
     !dataSourceView ||
diff --git a/types/src/front/api_handlers/public/vaults.ts b/types/src/front/api_handlers/public/vaults.ts
index 741326d4445c..7e7af46f8a1d 100644
--- a/types/src/front/api_handlers/public/vaults.ts
+++ b/types/src/front/api_handlers/public/vaults.ts
@@ -9,13 +9,13 @@ export const ContentSchema = t.type({
 
 export const PostDataSourceViewSchema = t.type({
   name: t.string,
-  parentsIn: t.array(t.string),
+  parentsIn: t.union([t.array(t.string), t.null]),
 });
 
 export type PostDataSourceViewType = t.TypeOf<typeof PostDataSourceViewSchema>;
 
 export const PatchDataSourceViewSchema = t.type({
-  parentsIn: t.array(t.string),
+  parentsIn: t.union([t.array(t.string), t.null]),
 });
 
 export type PatchDataSourceViewType = t.TypeOf<
@@ -41,11 +41,19 @@ export type PatchVaultRequestBodyType = t.TypeOf<
   typeof PostVaultRequestBodySchema
 >;
 
-export type ResourceCategory = "managed" | "files" | "webfolder" | "apps";
+export const RESOURCE_CATEGORIES = [
+  "managed",
+  "files",
+  "webfolder",
+  "apps",
+] as const;
+
+export type ResourceCategory = (typeof RESOURCE_CATEGORIES)[number];
 
 export type ResourceInfo = {
   createdAt: number;
   sId: string;
+  name: string;
   parentsIn?: string[] | null;
   connectorId?: string | null;
   connectorProvider?: ConnectorProvider | null;