From fd7dd193ed0e6853723f0ae55bbd20ab0aa23611 Mon Sep 17 00:00:00 2001
From: Andrew Wang <awwang3@uci.edu>
Date: Tue, 12 Mar 2024 14:53:07 -0700
Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20add=20restriction=20code=20?=
 =?UTF-8?q?filtering=20for=20Websoc=20lambda=20client?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/routes/v1/rest/websoc/+endpoint.ts    | 29 ++++++++++++++++++-
 apps/api/src/routes/v1/rest/websoc/lib.ts     |  1 -
 apps/api/src/routes/v1/rest/websoc/schema.ts  | 13 ++++++++-
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/apps/api/src/routes/v1/rest/websoc/+endpoint.ts b/apps/api/src/routes/v1/rest/websoc/+endpoint.ts
index f49ce8a1..e8a8f8aa 100644
--- a/apps/api/src/routes/v1/rest/websoc/+endpoint.ts
+++ b/apps/api/src/routes/v1/rest/websoc/+endpoint.ts
@@ -3,6 +3,7 @@ import { createHandler } from "@libs/lambda";
 import type { WebsocAPIResponse } from "@libs/uc-irvine-lib/websoc";
 import { notNull } from "@libs/utils";
 import { combineAndNormalizeResponses, sortResponse } from "@libs/websoc-utils";
+import type { z } from "zod";
 import { ZodError } from "zod";
 
 import { APILambdaClient } from "./APILambdaClient";
@@ -111,7 +112,9 @@ export const GET = createHandler(async (event, context, res) => {
       queries: normalizeQuery(parsedQuery),
     });
 
-    return res.createOKResult(websocResults, headers, requestId);
+    const filteredWebsocResults = filterResults(parsedQuery, websocResults);
+
+    return res.createOKResult(filteredWebsocResults, headers, requestId);
   } catch (error) {
     if (error instanceof ZodError) {
       const messages = error.issues.map((issue) => issue.message);
@@ -120,3 +123,27 @@ export const GET = createHandler(async (event, context, res) => {
     return res.createErrorResult(400, error, requestId);
   }
 }, onWarm);
+
+function filterResults(query: z.infer<typeof QuerySchema>, websocResults: WebsocAPIResponse) {
+  const excludeRestrictions = query.excludeRestrictionCodes ?? [];
+
+  if (excludeRestrictions.length) {
+    return websocResults.schools.map((school) => {
+      const filteredDepartments = school.departments.map((department) => {
+        const filteredCourses = department.courses.map((course) => {
+          const filteredSections = course.sections.filter(
+            (section) =>
+              !section.restrictions
+                .split(/ and | or /)
+                .some((code: string) => excludeRestrictions.includes(code)),
+          );
+          return { ...course, sections: filteredSections };
+        });
+        return { ...department, courses: filteredCourses };
+      });
+      return { ...school, departments: filteredDepartments };
+    });
+  }
+
+  return websocResults;
+}
diff --git a/apps/api/src/routes/v1/rest/websoc/lib.ts b/apps/api/src/routes/v1/rest/websoc/lib.ts
index 5ca473ab..13fd6dd8 100644
--- a/apps/api/src/routes/v1/rest/websoc/lib.ts
+++ b/apps/api/src/routes/v1/rest/websoc/lib.ts
@@ -202,7 +202,6 @@ export function constructPrismaQuery(parsedQuery: Query): Prisma.WebsocSectionWh
   }
 
   if (parsedQuery.excludeRestrictionCodes) {
-    console.log(parsedQuery.excludeRestrictionCodes as $Enums.RestrictionCode[]);
     AND.push({
       NOT: {
         restrictionCodes: {
diff --git a/apps/api/src/routes/v1/rest/websoc/schema.ts b/apps/api/src/routes/v1/rest/websoc/schema.ts
index 3bf5e404..05cd70d8 100644
--- a/apps/api/src/routes/v1/rest/websoc/schema.ts
+++ b/apps/api/src/routes/v1/rest/websoc/schema.ts
@@ -1,3 +1,4 @@
+import { $Enums } from "@libs/db";
 import {
   anyArray,
   cancelledCoursesOptions,
@@ -73,7 +74,17 @@ export const QuerySchema = z
   )
   .refine((x) => x.cacheOnly || x.building || !x.room, {
     message: 'If "building" is provided, "room" must also be provided',
-  });
+  })
+  .refine(
+    (x) =>
+      !x.excludeRestrictionCodes ||
+      x.excludeRestrictionCodes.every((code) =>
+        Object.values($Enums.RestrictionCode).includes(code as $Enums.RestrictionCode),
+      ),
+    {
+      message: `Restriction codes must be in [${Object.values($Enums.RestrictionCode).join(", ")}]`,
+    },
+  );
 
 /**
  * Type of the parsed query: useful for passing the query as input to other functions.