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 {
  fetchProfilesForSubjects,
  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("-t, --takendown", "Only sync takendown subjects")
  .option(
    "--reviewState <string>",
    "Items in specific review state to be synced"
  )
  .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;
      reviewState?: string;
      takendown?: boolean;
    }) => {
      if (!options.dataOnly) {
        const subjects = await fetchQueueItems(options);
        await saveQueueItems(subjects);
      }

      await fetchReposForSubjects();
      await fetchRecordsForSubjects();
      // await fetchProfilesForSubjects();
    }
  );

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(
      "--email <text/regex>",
      "Text or regex to be matched anywhere in the email address of the author."
    )
    .option(
      "--columns <string>",
      "Comma separated names of columns to include in the output"
    )
    .option(
      "--riskScore <number>",
      "Risk Score from threat signature filter",
      (value) => parseFloat(value)
    )
    .option("--lastCountry <string...>", "Last signin country of the user")
    .option("-td, --takendown", "Only use takendown subjects")
    .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")) {
        const mapped = subjects.map((sub) => {
          const { profile, ...rest } = sub;
          return {
            ...rest,
            name: sub.profile?.displayName || "",
          };
        });
        await writeCsvFile(options.file, mapped);
      } 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");
    });
  });