Bug fix to prevent queues from being exposed across environments.

Author: LaunaG

services/extract/src/pipeline/extract/tasks.py

@@ -2,6 +2,7 @@
 
 # Standard library imports
 import json
+import os
 import time
 from abc import ABC, abstractmethod
 
@@ -114,6 +115,14 @@ def __init__(self) -> None:
                 f"Django project not correctly configured. {e}"
             ) from None
 
+        # Parse environment variables
+        try:
+            self._env_flag = "-p-" if os.environ["ENV"] == "prod" else "-t-"
+        except KeyError as e:
+            raise RuntimeError(
+                f'Missing required environment variable "{e}". '
+            ) from None
+
         # List project queues
         try:
             self._queues = self.list_names()
@@ -202,7 +211,7 @@ def get_name(self, source: str) -> str:
             return matches[0]
 
     def list_names(self) -> list[str]:
-        """Fetches the names of all queues in the current project.
+        """Fetches the names of all queues in the current project environment.
 
         NOTE: Names are fully-qualified and have the format:
         `projects/PROJECT_ID/locations/LOCATION_ID/queues/QUEUE_ID`
@@ -239,7 +248,11 @@ def list_names(self) -> list[str]:
             )
 
         # Return names
-        return [queue["name"] for queue in r.json()["queues"]]
+        return [
+            queue["name"]
+            for queue in r.json()["queues"]
+            if self._env_flag in queue["name"]
+        ]
 
     def purge(self) -> None:
         """Purges all tasks from the configured queues.