From 7928ea2fff6e6d97df75b114c3acc1db31a2e7d7 Mon Sep 17 00:00:00 2001
From: Weves <chrisweaver101@gmail.com>
Date: Tue, 14 May 2024 15:40:13 -0700
Subject: [PATCH] Improve document set page UX

---
 backend/danswer/db/models.py                  | 12 ++++++++++++
 .../documents/sets/[documentSetId]/page.tsx   |  5 +++--
 web/src/app/admin/documents/sets/hooks.tsx    | 19 ++++++++++++++-----
 web/src/app/admin/documents/sets/new/page.tsx |  4 +++-
 4 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/backend/danswer/db/models.py b/backend/danswer/db/models.py
index 1991aee5a..37a5ded42 100644
--- a/backend/danswer/db/models.py
+++ b/backend/danswer/db/models.py
@@ -283,6 +283,10 @@ class ConnectorCredentialPair(Base):
     document_sets: Mapped[list["DocumentSet"]] = relationship(
         "DocumentSet",
         secondary=DocumentSet__ConnectorCredentialPair.__table__,
+        primaryjoin=(
+            (DocumentSet__ConnectorCredentialPair.connector_credential_pair_id == id)
+            & (DocumentSet__ConnectorCredentialPair.is_current.is_(True))
+        ),
         back_populates="connector_credential_pairs",
         overlaps="document_set",
     )
@@ -825,6 +829,14 @@ class DocumentSet(Base):
     connector_credential_pairs: Mapped[list[ConnectorCredentialPair]] = relationship(
         "ConnectorCredentialPair",
         secondary=DocumentSet__ConnectorCredentialPair.__table__,
+        primaryjoin=(
+            (DocumentSet__ConnectorCredentialPair.document_set_id == id)
+            & (DocumentSet__ConnectorCredentialPair.is_current.is_(True))
+        ),
+        secondaryjoin=(
+            DocumentSet__ConnectorCredentialPair.connector_credential_pair_id
+            == ConnectorCredentialPair.id
+        ),
         back_populates="document_sets",
         overlaps="document_set",
     )
diff --git a/web/src/app/admin/documents/sets/[documentSetId]/page.tsx b/web/src/app/admin/documents/sets/[documentSetId]/page.tsx
index 8efac3a11..960e11cc6 100644
--- a/web/src/app/admin/documents/sets/[documentSetId]/page.tsx
+++ b/web/src/app/admin/documents/sets/[documentSetId]/page.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import { ErrorCallout } from "@/components/ErrorCallout";
-import { useDocumentSets } from "../hooks";
+import { refreshDocumentSets, useDocumentSets } from "../hooks";
 import {
   useConnectorCredentialIndexingStatus,
   useUserGroups,
@@ -82,7 +82,8 @@ function Main({ documentSetId }: { documentSetId: number }) {
           ccPairs={ccPairs}
           userGroups={userGroups}
           onClose={() => {
-            router.push(`/admin/documents/sets?u=${Date.now()}`);
+            refreshDocumentSets();
+            router.push("/admin/documents/sets");
           }}
           setPopup={setPopup}
           existingDocumentSet={documentSet}
diff --git a/web/src/app/admin/documents/sets/hooks.tsx b/web/src/app/admin/documents/sets/hooks.tsx
index ee01f6f2f..fe7969ace 100644
--- a/web/src/app/admin/documents/sets/hooks.tsx
+++ b/web/src/app/admin/documents/sets/hooks.tsx
@@ -2,14 +2,23 @@ import { errorHandlingFetcher } from "@/lib/fetcher";
 import { DocumentSet } from "@/lib/types";
 import useSWR, { mutate } from "swr";
 
+const DOCUMENT_SETS_URL = "/api/manage/admin/document-set";
+
+export function refreshDocumentSets() {
+  mutate(DOCUMENT_SETS_URL);
+}
+
 export function useDocumentSets() {
-  const url = "/api/manage/admin/document-set";
-  const swrResponse = useSWR<DocumentSet[]>(url, errorHandlingFetcher, {
-    refreshInterval: 5000, // 5 seconds
-  });
+  const swrResponse = useSWR<DocumentSet[]>(
+    DOCUMENT_SETS_URL,
+    errorHandlingFetcher,
+    {
+      refreshInterval: 5000, // 5 seconds
+    }
+  );
 
   return {
     ...swrResponse,
-    refreshDocumentSets: () => mutate(url),
+    refreshDocumentSets: refreshDocumentSets,
   };
 }
diff --git a/web/src/app/admin/documents/sets/new/page.tsx b/web/src/app/admin/documents/sets/new/page.tsx
index c419acb65..24b337460 100644
--- a/web/src/app/admin/documents/sets/new/page.tsx
+++ b/web/src/app/admin/documents/sets/new/page.tsx
@@ -14,6 +14,7 @@ import { BackButton } from "@/components/BackButton";
 import { ErrorCallout } from "@/components/ErrorCallout";
 import { useRouter } from "next/navigation";
 import { UserGroup } from "@/lib/types";
+import { refreshDocumentSets } from "../hooks";
 
 function Main() {
   const { popup, setPopup } = usePopup();
@@ -50,7 +51,8 @@ function Main() {
           ccPairs={ccPairs}
           userGroups={userGroups}
           onClose={() => {
-            router.push(`/admin/documents/sets?u=${Date.now()}`);
+            refreshDocumentSets();
+            router.push("/admin/documents/sets");
           }}
           setPopup={setPopup}
         />