minimize calls to firebase for job status data

src/App.tsx

@@ -1,10 +1,10 @@
 import { useState } from "react";
 import { v4 as uuidv4 } from 'uuid';
 import { Layout, Typography } from "antd";
-import { getResultPath, getDocById, getJobStatus, addRecipe } from "./utils/firebase";
+import { getJobStatus, addRecipe } from "./utils/firebase";
 import { getFirebaseRecipe, jsonToString } from "./utils/recipeLoader";
 import { getSubmitPackingUrl, JOB_STATUS } from "./constants/aws";
-import { FIRESTORE_COLLECTIONS, FIRESTORE_FIELDS } from "./constants/firebase";
+import { FIRESTORE_FIELDS } from "./constants/firebase";
 import { SIMULARIUM_EMBED_URL } from "./constants/urls";
 import PackingInput from "./components/PackingInput";
 import Viewer from "./components/Viewer";
@@ -20,6 +20,7 @@ function App() {
     const [jobStatus, setJobStatus] = useState("");
     const [jobLogs, setJobLogs] = useState<string>("");
     const [resultUrl, setResultUrl] = useState<string>("");
+    const [outputDir, setOutputDir] = useState<string>("");
     const [runTime, setRunTime] = useState<number>(0);
 
     let start = 0;
@@ -88,34 +89,24 @@ function App() {
     const checkStatus = async (jobIdFromSubmit: string) => {
         const id = jobIdFromSubmit || jobId;
         let localJobStatus = await getJobStatus(id);
-        while (localJobStatus !== JOB_STATUS.DONE && localJobStatus !== JOB_STATUS.FAILED) {
+        while (localJobStatus?.status !== JOB_STATUS.DONE && localJobStatus?.status !== JOB_STATUS.FAILED) {
             await sleep(500);
             const newJobStatus = await getJobStatus(id);
-            if (localJobStatus !== newJobStatus) {
+            if (newJobStatus && localJobStatus?.status !== newJobStatus.status) {
                 localJobStatus = newJobStatus;
-                setJobStatus(newJobStatus);
+                setJobStatus(newJobStatus.status);
             }
         }
         const range = (Date.now() - start) / 1000;
         setRunTime(range);
-        if (localJobStatus == JOB_STATUS.DONE) {
-            fetchResultUrl(id);
-        } else if (localJobStatus == JOB_STATUS.FAILED) {
-            getLogs(id);
+        if (localJobStatus.status == JOB_STATUS.DONE) {
+            setResultUrl(SIMULARIUM_EMBED_URL + localJobStatus.result_path);
+            setOutputDir(localJobStatus.outputs_directory);
+        } else if (localJobStatus.status == JOB_STATUS.FAILED) {
+            setJobLogs(localJobStatus.error_message);
         }
     };
 
-    const fetchResultUrl = async (jobIdFromSubmit?: string) => {
-        const id = jobIdFromSubmit || jobId;
-        const url = await getResultPath(id);
-        setResultUrl(SIMULARIUM_EMBED_URL + url);
-    };
-
-    const getLogs = async (jobIdFromSubmit?: string) => {
-        const id = jobIdFromSubmit || jobId;
-        const logStr: string = await getDocById(FIRESTORE_COLLECTIONS.JOB_STATUS, id);
-        setJobLogs(logStr);
-    };
     const showLogs = jobStatus == JOB_STATUS.FAILED;
 
     return (
@@ -126,8 +117,8 @@ function App() {
             </Header>
             <Content className="content-container">
                 <PackingInput startPacking={startPacking} />
-                {jobStatus && <StatusBar jobStatus={jobStatus} runTime={runTime} jobId={jobId} />}
-                {showLogs && <ErrorLogs errorLogs={jobLogs} getLogs={getLogs} />}
+                {jobStatus && <StatusBar jobStatus={jobStatus} runTime={runTime} jobId={jobId} outputDir={outputDir} />}
+                {showLogs && <ErrorLogs errorLogs={jobLogs} />}
             </Content>
             {resultUrl && <Viewer resultUrl={resultUrl} />}
         </div>

src/components/ErrorLogs/index.tsx

@@ -4,19 +4,14 @@ import "./style.css";
 
 interface ErrorLogsProps {
     errorLogs: string;
-    getLogs: () => Promise<void>;
 }
 
 const ErrorLogs = (props: ErrorLogsProps): JSX.Element => {
-    const { errorLogs, getLogs } = props;
+    const { errorLogs } = props;
     const [viewErrorLogs, setViewErrorLogs] = useState<boolean>(true);
 
-    const toggleLogs = async () => {
-        if (errorLogs.length === 0) {
-            await getLogs();
-        } else {
-            setViewErrorLogs(!viewErrorLogs);
-        }
+    const toggleLogs = () => {
+        setViewErrorLogs(!viewErrorLogs);
     }
     const items = [{
         key: "1",

src/components/StatusBar/index.tsx

@@ -8,15 +8,16 @@ interface StatusBarProps {
     jobStatus: string;
     runTime: number;
     jobId: string;
+    outputDir: string;
 } 
 
 const StatusBar = (props: StatusBarProps): JSX.Element => {
-    const { jobStatus, runTime, jobId } = props;
+    const { jobStatus, runTime, jobId, outputDir } = props;
     const [isDownloading, setIsDownloading] = useState(false);
 
     const downloadResults = async (jobId: string) => {
         setIsDownloading(true);
-        await downloadOutputs(jobId);
+        await downloadOutputs(jobId, outputDir);
         setIsDownloading(false);
     }
 

src/types/index.ts

@@ -5,40 +5,28 @@ export interface Document {
     recipe?: string;
     config?: string;
     editable_fields?: string[];
-}
+};
 
 export type FirestoreDoc = Document & {
     id: string;
 };
 
-export interface AWSBatchJobsResponse {
-    jobs: Array<{
-        status: string;
-        container: {
-            logStreamName: string;
-        };
-    }>;
-}
-
-export interface CloudWatchLogsResponse {
-    events: Array<{
-        message: string;
-    }>;
-}
-
-export type FirebaseDict = {
-    [key: string]: Dictionary<string>;
-};
-
 export interface Dictionary<T> {
     [Key: string]: T;
-}
+};
 
 export type PackingInputs = {
     config: string;
     recipe: string;
     editable_fields?: EditableField[];
-}
+};
+
+export type JobStatusObject = {
+    status: string;
+    error_message: string;
+    outputs_directory: string;
+    result_path: string;
+};
 
 export type EditableField = {
     id: string;
@@ -67,14 +55,14 @@ export type GradientOption = {
     packing_mode_path?: string;
     strength_description?: string;
     strength_display_name?: string;
-}
+};
 
 export interface RefsByCollection {
     recipes: Dictionary<FirebaseRecipe>;
     composition: Dictionary<FirebaseComposition>;
     objects: Dictionary<FirebaseObject>;
     gradients: Dictionary<FirebaseGradient>;
-}
+};
 
 export interface FirebaseObject {
     name: string;
@@ -110,7 +98,7 @@ export interface FirebaseObject {
     rejection_threshold?: number;
     rotation_axis?: Array<number>;
     rotation_range?: number;
-}
+};
 
 export interface FirebaseGradient {
     name: string;
@@ -124,12 +112,12 @@ export interface FirebaseGradient {
     pick_mode?: string;
     reversed?: boolean;
     invert?: string;
-}
+};
 
 export type RegionObject = {
     count: number;
     object: string;
-}
+};
 
 export interface FirebaseComposition {
     name: string;
@@ -177,4 +165,4 @@ export type ViewableRecipe = {
     composition?: Dictionary<ViewableComposition>;
     objects?: Dictionary<ViewableObject>;
     gradients?: Dictionary<ViewableGradient>;
-}
+};

src/utils/aws.ts

@@ -78,7 +78,11 @@ export const downloadOutputsFromS3 = async (outputsDir: string, jobId: string) =
     console.log(`Downloaded zip with ${filesAdded} files`);
 };
 
-export const downloadOutputs = async (jobId: string) => {
-    const outputsDir = await getOutputsDirectory(jobId);
+export const downloadOutputs = async (jobId: string, outputsDir?: string) => {
+    if (!outputsDir) {
+        // If uploading result files took too long, outputDir may not have been
+        // set when the job completed. Fetch it from Firestore.
+        outputsDir = await getOutputsDirectory(jobId);
+    }
     await downloadOutputsFromS3(outputsDir, jobId);
 }

src/utils/firebase.ts

@@ -24,6 +24,7 @@ import {
     PackingInputs,
     Dictionary,
     EditableField,
+    JobStatusObject,
 } from "../types";
 
 const getEnvVar = (key: string): string => {
@@ -69,14 +70,6 @@ const queryDocumentsByIds = async (collectionName: string, ids: string[]) => {
     return await getDocs(q);
 };
 
-const queryDocumentsByField = async (collectionName: string, field: string, value: string | number | boolean) => {
-    const q = query(
-        collection(db, collectionName),
-        where(field, "==", value)
-    );
-    return await getDocs(q);
-};
-
 const queryAllDocuments = async (collectionName: string) => {
     const q = query(collection(db, collectionName));
     return await getDocs(q);
@@ -106,15 +99,16 @@ const extractSingleDocumentData = (querySnapshot: QuerySnapshot<DocumentData>, f
 };
 
 // Query functions for our use case using generic functions
-const getResultPath = async (jobId: string) => {
-    const querySnapshot = await queryDocumentsByField(FIRESTORE_COLLECTIONS.RESULTS, FIRESTORE_FIELDS.BATCH_JOB_ID, jobId);
-    return extractSingleDocumentData(querySnapshot, FIRESTORE_FIELDS.URL);
-};
-
-const getJobStatus = async (jobId: string) => {
+const getJobStatus = async (jobId: string): Promise<JobStatusObject|undefined> => {
     const querySnapshot = await queryDocumentById(FIRESTORE_COLLECTIONS.JOB_STATUS, jobId);
-    return extractSingleDocumentData(querySnapshot, FIRESTORE_FIELDS.STATUS);
-}
+    const docs = querySnapshot.docs.map((doc) => ({
+        status: doc.data().status,
+        error_message: doc.data().error_message,
+        outputs_directory: doc.data().outputs_directory,
+        result_path: doc.data().result_path,
+    }));
+    return docs[0] || undefined;
+};
 
 const getOutputsDirectory = async (jobId: string) => {
     const querySnapshot = await queryDocumentById(FIRESTORE_COLLECTIONS.JOB_STATUS, jobId);
@@ -167,12 +161,6 @@ const getPackingInputsDict = async (): Promise<Dictionary<PackingInputs>> => {
     return inputsDict;
 }
 
-const getDocById = async (coll: string, id: string) => {
-    const docs = await getAllDocsFromCollection(coll);
-    const doc = docs.find(d => d.id === id);
-    return JSON.stringify(doc, null, 2);
-}
-
 const getDocsByIds = async (coll: string, ids: string[]) => {
     const querySnapshot = await queryDocumentsByIds(coll, ids);
     const docs = querySnapshot.docs.map((doc) => ({
@@ -214,4 +202,4 @@ const docCleanup = async () => {
         console.log(`Cleaned up ${deletePromises.length} documents from ${collectionConfig.name}`);
     }
 }
-export { db, queryDocumentById, getDocById, getDocsByIds, getJobStatus, getResultPath, addRecipe, docCleanup, getPackingInputsDict, getOutputsDirectory };
+export { db, queryDocumentById, getDocsByIds, getJobStatus, addRecipe, docCleanup, getPackingInputsDict, getOutputsDirectory };