auron/src/commands/queue.ts

Interact with your ATProto labeler from your terminal

import { Command } from "commander";
import { database } from "../services/db";
import { withLoader } from "../utils/loader";
import {
  fetchQueueItems,
  fetchRecordsForSubjects,
  fetchReposForSubjects,
  processSubjects,
  saveQueueItems,
} from "../controllers/queue";
import { writeCsvFile, writeJsonFile } from "../services/file";

export const queueCommand = new Command("queue");

queueCommand
  .command("sync")
  .option(
    "-n, --count <number>",
    "Max number of items from the queue to sync",
    (value) => parseInt(value, 10),
    100
  )
  .option("--cursor <string>", "Cursor to start fetching from")
  .option(
    "--data-only",
    "Only fetch record and repo data, without syncing queue items"
  )
  .description("Fetch your ozone queue and store it in the local database")
  .action(
    async (options: { count: number; cursor?: string; dataOnly?: boolean }) => {
      if (!options.dataOnly) {
        const subjects = await fetchQueueItems(options);
        await saveQueueItems(subjects);
      }
      await fetchReposForSubjects();
      await fetchRecordsForSubjects();
    }
  );

const addSubjectListOptions = (command: Command): Command => {
  return command
    .option(
      "-n, --count <number>",
      "Max number of items from the queue",
      (value) => parseInt(value, 10)
    )
    .option(
      "-t, --type <type>",
      "Subject type filter, account or record",
      (value) => {
        if (value && value !== "account" && value !== "record") {
          throw new Error(
            'Invalid type. Allowed values are "account" or "record".'
          );
        }
        return value;
      }
    )
    .option(
      "--bio <keyword/text>",
      "Keyword to be matched in profile bio of subject"
    )
    .option(
      "--keyword <keyword/text>",
      "Keyword to be matched anywhere in the profile bio/record content etc."
    )
    .option("--cursor <string>", "Cursor to start fetching from");
};

addSubjectListOptions(queueCommand.command("export"))
  .description("Export queue from your local database")
  .requiredOption("-f, --file <file>", "Export file path") // Make the file option required
  .action(async (options) => {
    withLoader(`Exporting queue items...`, async (updateMessage) => {
      const subjects = await processSubjects(options);

      if (subjects.length === 0) {
        updateMessage(`No queue items found`);
        return;
      }

      updateMessage(
        `Found ${subjects.length} queue items. Exporting to ${options.file}`
      );
      if (options.file.endsWith(".csv")) {
        await writeCsvFile(options.file, subjects);
      } else {
        await writeJsonFile(options.file, subjects);
      }

      const lastItem = subjects[subjects.length - 1];
      updateMessage(
        `${subjects.length} queue items exported to ${options.file}. Cursor: ${lastItem.lastReportedAt}`
      );
    });
  });

addSubjectListOptions(queueCommand.command("pipe"))
  .description("Output 1 item at a time from your queue to stdout")
  .action(async (options) => {
    const subjects = await processSubjects(options);

    if (subjects.length === 0) {
      return;
    }

    for (const subject of subjects) {
      console.log(JSON.stringify(subject));
    }
  });

queueCommand
  .command("clear")
  .description("Clear your locally stored queue")
  .action(async () => {
    withLoader(`Clearing queue items...`, async (logMessage) => {
      await database.clearSubjects();
      logMessage("Queue items cleared from the database");
    });
  });