From 4d6fa552f36a0348ca60b60335525038b2cbf5d7 Mon Sep 17 00:00:00 2001
From: Luke Chavers <lucas.chavers@gsa.gov>
Date: Fri, 18 Oct 2024 17:40:20 -0400
Subject: [PATCH] Create cli function to enqueue a limited set of websites

---
 .github/workflows/enqueue-limited-scans.yml | 34 +++++++++++++++++++++
 apps/cli/src/main.ts                        | 24 +++++++++++++++
 apps/cli/src/queue.controller.ts            |  8 +++--
 3 files changed, 64 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/enqueue-limited-scans.yml

diff --git a/.github/workflows/enqueue-limited-scans.yml b/.github/workflows/enqueue-limited-scans.yml
new file mode 100644
index 00000000..213226ef
--- /dev/null
+++ b/.github/workflows/enqueue-limited-scans.yml
@@ -0,0 +1,34 @@
+---
+name: Enqueue limited scans
+
+# yamllint disable-line rule:truthy
+on:
+  workflow_dispatch:
+    inputs:
+      cf_space:
+        description: deployed cloud.gov space
+        required: true
+        default: prod
+      site_limit:
+        description: number of sites to queue
+        required: true
+        default: 10
+        type: number
+
+jobs:
+  enqueue:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: run enqueue-limited-scans cli job in chosen space
+        uses: cloud-gov/cg-cli-tools@main
+        with:
+          cf_api: https://api.fr.cloud.gov
+          cf_username: ${{ secrets.CF_USERNAME }}
+          cf_password: ${{ secrets.CF_PASSWORD }}
+          cf_org: gsatts-sitescan
+          cf_space: ${{ github.event.inputs.cf_space || 'prod' }}
+          cf_command: run-task site-scanner-consumer
+            -c "node dist/apps/cli/main.js enqueue-limited-scans --limit ${{ github.event.inputs.site_limit }}"
+            -k 2G -m 2G
+            --name github-action-enqueue-limited-scans-${{ github.run_id }}
diff --git a/apps/cli/src/main.ts b/apps/cli/src/main.ts
index 00a5b215..64923e1e 100644
--- a/apps/cli/src/main.ts
+++ b/apps/cli/src/main.ts
@@ -75,6 +75,17 @@ async function enqueueScans() {
   await nestApp.close();
 }
 
+async function enqueueLimitedScans(cmdObj) {
+  const nestApp = await bootstrap();
+  const logger = createCommandLogger('enqueue-limited-scans', { limit: cmdObj.limit });
+  const controller = nestApp.get(QueueController);
+  logger.info('enqueueing limited scan jobs');
+
+  await controller.queueScans(cmdObj.limit);
+  printMemoryUsage(logger);
+  await nestApp.close();
+}
+
 async function createSnapshot() {
   const nestApp = await bootstrap();
   const logger = createCommandLogger('create-snapshot');
@@ -168,6 +179,19 @@ async function main() {
       'enqueue-scans adds each target in the Website database table to the redis queue',
     )
     .action(enqueueScans);
+  
+  // queue-limited-scans
+  program
+    .command('enqueue-limited-scans')
+    .description(
+      'enqueue-scans-limit adds a limited set of targets from the Website database table to the redis queue',
+    )
+    .option(
+      '--limit <number>',
+      'limits the enqueue service to <number> urls',
+      parseInt,
+    )
+    .action((cmdObj) => enqueueLimitedScans(cmdObj));
 
   // create-snapshot
   program
diff --git a/apps/cli/src/queue.controller.ts b/apps/cli/src/queue.controller.ts
index 3d8f6216..72fcb1ae 100644
--- a/apps/cli/src/queue.controller.ts
+++ b/apps/cli/src/queue.controller.ts
@@ -14,11 +14,15 @@ export class QueueController {
     private websiteService: WebsiteService,
   ) {}
 
-  async queueScans() {
+  async queueScans(limit?: number) {
     this.logger.log('starting to enqueue scans...');
 
     try {
-      const websites = await this.websiteService.findAllWebsites();
+      let websites = await this.websiteService.findAllWebsites();
+
+      if(limit) {
+        websites = websites.slice(0, limit);
+      }
 
       this.logger.log(`adding ${websites.length} websites to the queue`);