Implement take step UI for query builder

frontend/src/components/query/query-builder/query-builder.tsx

@@ -3,10 +3,11 @@ import {
   Pipeline,
   SelectStep,
   AggregateStep,
+  RelateStep,
+  TakeStep,
   Step,
   StepIdentifier,
   StepIdentifierEnum,
-  RelateStep,
 } from "@/definitions/pipeline";
 import { Button, NonIdealState, Section } from "@blueprintjs/core";
 import Loading from "@/app/loading";
@@ -16,6 +17,7 @@ import FromStepComponent from "./steps/from-step";
 import SelectStepComponent from "./steps/select-step";
 import AggregateStepComponent from "./steps/aggregate-step";
 import RelateStepComponent from "./steps/relate/relate-step";
+import TakeStepComponent from "./steps/take-step";
 import { usePipelineSchema } from "@/data/use-user-query";
 import { useState, useEffect } from "react";
 import * as _ from "lodash";
@@ -130,10 +132,22 @@ export default function QueryBuilder({
             create={create}
           />
         );
+      case StepIdentifierEnum.Take:
+        return (
+          <TakeStepComponent
+            index={index}
+            step={!!step ? (step as TakeStep) : step}
+            pipeline={pipeline}
+            setPipeline={setPipeline}
+            edit={editStepIndex === index}
+            setEditStepIndex={setEditStepIndex}
+            setNewStepType={setNewStepType}
+            create={create}
+          />
+        );
       case StepIdentifierEnum.Derive:
       case StepIdentifierEnum.Filter:
       case StepIdentifierEnum.Order:
-      case StepIdentifierEnum.Take:
         return (
           <Section className="flex-none w-full py-2 rounded-sm">
             <NonIdealState

frontend/src/components/query/query-builder/steps/take-step.tsx

@@ -0,0 +1,292 @@
+"use client";
+import {
+  Pipeline,
+  TakeStep,
+  Step,
+  StepIdentifierEnum,
+} from "@/definitions/pipeline";
+import {
+  Button,
+  ButtonGroup,
+  Divider,
+  Menu,
+  MenuItem,
+  NumericInput,
+  Popover,
+  Section,
+  Text,
+} from "@blueprintjs/core";
+import { NewStepSelection } from "../query-builder";
+import Loading from "@/app/loading";
+import { ErrorDisplay } from "@/components/error-display";
+import InvalidStepPopover from "../invalid-step-popover";
+import { useEffect } from "react";
+import { useField } from "@/utils/use-field";
+import { usePipelineSchema } from "@/data/use-user-query";
+import * as _ from "lodash";
+
+interface TakeStepProps {
+  index: number;
+  step: TakeStep | null;
+  pipeline: Pipeline;
+  setPipeline: (value: Pipeline) => void;
+  edit: boolean;
+  setEditStepIndex: (value: number | null) => void;
+  setNewStepType: (value: NewStepSelection | null) => void;
+  create?: boolean;
+}
+
+export default function TakeStepComponent({
+  index,
+  step,
+  pipeline,
+  setPipeline,
+  edit,
+  setEditStepIndex,
+  setNewStepType,
+  create,
+}: TakeStepProps) {
+  const offset = useField<number | undefined>(undefined);
+  const limit = useField<number | undefined>(undefined, {
+    valueTransformer: (value) => {
+      return value === 0 ? 1 : value;
+    },
+  });
+
+  const {
+    data: inputSchema,
+    isLoading: isLoadingInputSchema,
+    error: inputSchemaError,
+  } = usePipelineSchema({
+    ...pipeline,
+    steps: _.slice(pipeline.steps, 0, index),
+  });
+
+  const {
+    data: schema,
+    isLoading: isLoadingSchema,
+    error: schemaError,
+  } = usePipelineSchema({
+    ...pipeline,
+    steps: _.slice(pipeline.steps, 0, index + 1),
+  });
+
+  useEffect(() => {
+    resetFields();
+  }, [step]);
+
+  function resetFields() {
+    if (step) {
+      offset.onValueChange(step.offset);
+      limit.onValueChange(step.limit);
+    } else {
+      offset.onValueChange(undefined);
+      limit.onValueChange(undefined);
+    }
+  }
+
+  function getAdditionalClasses() {
+    if (inputSchema && !inputSchema.success) {
+      return "border-2 border-gold";
+    } else if (schema && !schema.success) {
+      return "border-2 border-error";
+    }
+  }
+
+  function renderTitle() {
+    // If we are creating a new step, editing a step, or the step is not defined, show the default title
+    // Also show the default title if the schemas are loading or errored
+    if (
+      edit ||
+      create ||
+      !step ||
+      isLoadingSchema ||
+      isLoadingInputSchema ||
+      inputSchemaError ||
+      schemaError
+    ) {
+      return <Text className="text-xl grow-0">Take:</Text>;
+    } else if (inputSchema && !inputSchema.success) {
+      return (
+        <div className="flex flex-row items-center">
+          <Text className="text-xl grow-0">Take:</Text>
+          <ButtonGroup
+            minimal={true}
+            vertical={false}
+            className="flex flex-row items-center gap-2 ml-3"
+          >
+            <div className="flex flex-row items-center">
+              <Text className="text-base font-normal">limit</Text>
+              <Text className="ml-2 font-bold">{step.limit}</Text>
+            </div>
+            <Divider className="h-3.5" />
+            <div className="flex flex-row items-center">
+              <Text className="text-base font-normal">offset</Text>
+              <Text className="ml-2 font-bold">{step.offset}</Text>
+            </div>
+          </ButtonGroup>
+        </div>
+      );
+    } else if (schema && !schema.success) {
+      return (
+        <div className="flex flex-row items-center">
+          <Text className="text-xl grow-0">Take:</Text>
+          <InvalidStepPopover errors={schema!.error.issues} />
+        </div>
+      );
+    } else {
+      return (
+        <div className="flex flex-row items-center">
+          <Text className="text-xl grow-0">Take:</Text>
+          <ButtonGroup
+            minimal={true}
+            vertical={false}
+            className="flex flex-row items-center gap-2 ml-3"
+          >
+            <div className="flex flex-row items-center">
+              <Text className="text-base font-normal">limit</Text>
+              <Text className="ml-2 font-bold">{step.limit}</Text>
+            </div>
+            <Divider className="h-3.5" />
+            <div className="flex flex-row items-center">
+              <Text className="text-base font-normal">offset</Text>
+              <Text className="ml-2 font-bold">{step.offset}</Text>
+            </div>
+          </ButtonGroup>
+        </div>
+      );
+    }
+  }
+
+  function renderRightElement() {
+    if (create) {
+      return (
+        <>
+          <Button
+            alignText="left"
+            disabled={limit.value === undefined || offset.value === undefined}
+            text="Add step"
+            onClick={() => {
+              setNewStepType(null);
+              const newStep = {
+                type: StepIdentifierEnum.Take,
+                limit: limit.value,
+                offset: offset.value,
+              } as TakeStep;
+              const newSteps: Step[] = [...pipeline.steps];
+              newSteps.splice(index, 0, newStep as Step);
+              setPipeline({ ...pipeline, steps: newSteps });
+            }}
+          />
+          <Button
+            className="ml-2"
+            alignText="left"
+            text="Cancel step"
+            onClick={() => {
+              setNewStepType(null);
+            }}
+          />
+        </>
+      );
+    } else {
+      if (edit) {
+        return (
+          <>
+            <Button
+              alignText="left"
+              disabled={limit.value === undefined || offset.value === undefined}
+              text="Confirm step"
+              onClick={() => {
+                const updatedStep = {
+                  type: StepIdentifierEnum.Take,
+                  limit: limit.value,
+                  offset: offset.value,
+                } as TakeStep;
+                const newSteps: Step[] = [...pipeline.steps];
+                newSteps.splice(index, 1, updatedStep as Step);
+                setPipeline({ ...pipeline, steps: newSteps });
+              }}
+            />
+            <Button
+              className="ml-2"
+              alignText="left"
+              text="Cancel"
+              onClick={() => {
+                resetFields();
+                setEditStepIndex(null);
+              }}
+            />
+          </>
+        );
+      } else {
+        return (
+          <Popover
+            content={
+              <Menu>
+                <MenuItem
+                  icon="edit"
+                  text="Edit step"
+                  disabled={!!inputSchema && !inputSchema.success}
+                  onClick={() => setEditStepIndex(index)}
+                />
+                <MenuItem
+                  icon="trash"
+                  text="Delete step"
+                  onClick={() => {
+                    const newSteps: Step[] = [...pipeline.steps];
+                    newSteps.splice(index, 1);
+                    setPipeline({ ...pipeline, steps: newSteps });
+                  }}
+                />
+              </Menu>
+            }
+            placement="bottom"
+          >
+            <Button alignText="left" rightIcon="caret-down" text="Options" />
+          </Popover>
+        );
+      }
+    }
+  }
+
+  function renderContent() {
+    if (isLoadingSchema || isLoadingInputSchema) {
+      return (
+        <div className="flex flex-row justify-center my-1 h-fit">
+          <Loading />
+        </div>
+      );
+    } else if (schemaError || inputSchemaError) {
+      return <ErrorDisplay description={schemaError || inputSchemaError} />;
+    } else if (!inputSchema!.success) {
+      return null;
+    } else if (edit || create || !step) {
+      return (
+        <div className="flex flex-row gap-3 mx-3 my-2">
+          <NumericInput
+            placeholder="limit"
+            allowNumericCharactersOnly={true}
+            min={1}
+            {...limit}
+          />
+          <NumericInput
+            placeholder="offset"
+            allowNumericCharactersOnly={true}
+            min={0}
+            {...offset}
+          />
+        </div>
+      );
+    }
+  }
+
+  return (
+    <Section
+      className={`flex-none w-full rounded-sm ${getAdditionalClasses()}`}
+      title={renderTitle()}
+      rightElement={<div className="flex flex-row">{renderRightElement()}</div>}
+    >
+      {renderContent()}
+    </Section>
+  );
+}

frontend/src/utils/pipeline/validate-pipeline-step.tsx

@@ -72,9 +72,10 @@ function getPipelineStepValidator(
         tables,
         relations,
       );
-    case StepIdentifierEnum.Filter:
-    case StepIdentifierEnum.Order:
     case StepIdentifierEnum.Take:
+      return createTakeStepValidator(stepIndex);
+    case StepIdentifierEnum.Order:
+    case StepIdentifierEnum.Filter:
     case StepIdentifierEnum.Derive:
     default:
       throw new Error("Invalid step type");
@@ -308,3 +309,29 @@ function createRelateStepValidator(
   );
   return relateValidator;
 }
+
+function createTakeStepValidator(stepIndex: number) {
+  const takeValidator = TakeStepSchema.superRefine(
+    (step: TakeStep, ctx: any) => {
+      if (typeof step.limit !== "number") {
+        ctx.addIssue({
+          code: z.ZodIssueCode.custom,
+          message: `Limit must be type 'number', recieved type '${typeof step.limit}'`,
+          path: [stepIndex.toString(), `step ${stepIndex + 1} - Take`, "limit"],
+        });
+      }
+      if (typeof step.offset !== "number") {
+        ctx.addIssue({
+          code: z.ZodIssueCode.custom,
+          message: `Offset must be type 'number', recieved type '${typeof step.offset}'`,
+          path: [
+            stepIndex.toString(),
+            `step ${stepIndex + 1} - Take`,
+            "offset",
+          ],
+        });
+      }
+    },
+  );
+  return takeValidator;
+}