auron/src/controllers/queue.ts

Interact with your ATProto labeler from your terminal

import EventEmitter from "events";
import { withLoader } from "../utils/loader";
import { getQueueItems, getRecords, getRepos } from "../api/atproto";
import { database } from "../services/db";
import { ToolsOzoneModerationDefs } from "@atproto/api";
import {
  transformRecordViewToRecord,
  transformRepoViewToRepo,
  transformStatusToSubject,
} from "../services/subject/transformers";
import { chunkArray } from "@atproto/common";

export const fetchQueueItems = async (options: {
  count: number;
  cursor?: string;
}) => {
  const emitter = new EventEmitter();
  const subjects = await withLoader(
    `Fetching queue items...`,
    async (updateMessage) => {
      emitter.on("update", ({ nextCursor, subjectCount, maxCount }) => {
        updateMessage(
          `Fetched page with cursor: ${nextCursor}, total subjects: ${subjectCount} out of ${maxCount}`
        );
      });

      const { cursor, subjectStatuses } = await getQueueItems(
        {
          maxCount: options.count ? Number(options.count) : undefined,
          cursor: options.cursor,
        },
        emitter
      );
      updateMessage(
        `Fetched ${subjectStatuses.length} queue items. Cursor: ${cursor}`
      );
      return subjectStatuses;
    }
  );

  return subjects;
};

export const saveQueueItems = async (
  subjects: ToolsOzoneModerationDefs.SubjectStatusView[]
) => {
  await withLoader(
    `Storing ${subjects.length} queue items...`,
    async (updateMessage) => {
      if (subjects.length === 0) {
        updateMessage("No queue items to store");
        throw new Error("No queue items to store");
      }
      for (const chunk of chunkArray(subjects, 500)) {
        await database.insertSubjects(chunk.map(transformStatusToSubject));
        updateMessage(
          `Stored ${chunk.length} out of ${subjects.length} queue items in local database`
        );
      }
      updateMessage(`Stored ${subjects.length} queue items in local database`);
    }
  );
};

export const fetchReposForSubjects = async () => {
  const emitter = new EventEmitter();
  const subjects = await withLoader(
    `Fetching repos for subjects...`,
    async (updateMessage) => {
      emitter.on("update", ({ total, repoCount }) => {
        updateMessage(`Fetched ${repoCount} repos out of ${total}`);
      });

      const dids = await database.getMissingRepoDids();
      updateMessage(`Found ${dids.length} missing repos to be fetched`);

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

      const repos = await getRepos({ dids }, emitter);
      updateMessage(
        `Fetched ${repos.length} repos. Saving in local database now`
      );

      for (const chunk of chunkArray(repos, 500)) {
        await database.saveRepos(chunk.map(transformRepoViewToRepo));
        updateMessage(
          `Saved ${chunk.length} out of ${repos.length} repos in local database now`
        );
      }

      updateMessage(`Saved ${repos.length} repos in local database`);
      return;
    }
  );

  return subjects;
};

export const fetchRecordsForSubjects = async () => {
  const emitter = new EventEmitter();
  const subjects = await withLoader(
    `Fetching records for subjects...`,
    async (updateMessage) => {
      emitter.on("update", ({ total, recordCount }) => {
        updateMessage(`Fetched ${recordCount} records out of ${total}`);
      });

      const uris = await database.getMissingRecordUris();
      updateMessage(`Found ${uris.length} missing records to be fetched`);

      if (uris.length === 0) {
        return;
      }
      const records = await getRecords({ uris }, emitter);
      updateMessage(
        `Fetched ${records.length} records. Saving in local database now`
      );

      for (const chunk of chunkArray(records, 500)) {
        await database.saveRecords(chunk.map(transformRecordViewToRecord));
        updateMessage(
          `Saved ${chunk.length} out of ${records.length} records in local database`
        );
      }

      updateMessage(`Saved ${records.length} records in local database`);
      return;
    }
  );

  return subjects;
};

export const processSubjects = async (options: {
  type?: string;
  cursor?: string;
  count?: number;
  bio?: string;
  keyword?: string;
}) => {
  let subjects = await database.listSubjects({
    subjectType: options.type,
    cursor: options.cursor,
    limit: options.count,
  });

  if (options.bio || options.keyword) {
    const keyword = `${options.bio || options.keyword}`.toLowerCase();
    subjects = subjects.filter((subject) => {
      const profile = JSON.parse(subject.profile || "{}");
      return profile.description?.toLowerCase().includes(keyword);
    });
  }

  if (options.keyword) {
    const keyword = options.keyword.toLowerCase();
    subjects = subjects.filter((subject) => {
      const recordValue = JSON.parse(subject.value || "{}");
      return recordValue.text?.toLowerCase().includes(keyword);
    });
  }

  return subjects;
};