Migrate to RabbitMQ

CONTRIBUTING.md

@@ -137,7 +137,6 @@ You'll need to have the following environment variables set.
 | `AUTH0_DOMAIN` | str | "vipyrsec.us.auth0.com" | Authentication domain for Auth0 |
 | `AUTH0_AUDIENCE` | str | "dragonfly.vipyrsec.com" | Audience field for Auth0 |
 | `DRAGONFLY_GITHUB_TOKEN` | str | | Github PAT for accessing YARA rules in the security-intelligence repository |
-| `JOB_TIMEOUT` | int | 60 \* 2 | The maximum time to wait for clients to respond with job results. After this time has elapsed, the server will begin distributing this job to other clients |
 | | | | |
 | `EMAIL_RECIPIENT` | str | "[email protected]" | The recipient address of report emails |
 | `BCC_RECIPIENTS` | set | set() | Additional addresses that should be BCC'd in email reports. Defaults to an empty set. |

alembic/versions/58fcd14830ae_drop_unneeded_columns.py

@@ -0,0 +1,61 @@
+"""Drop unneeded columns
+
+Revision ID: 58fcd14830ae
+Revises: 9d70bac89aa5
+Create Date: 2023-08-13 16:18:48.445009
+
+"""
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "58fcd14830ae"
+down_revision = "94f1cd79b7d4"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.drop_table("download_urls")
+ op.drop_column("scans", "fail_reason")
+ op.drop_column("scans", "pending_by")
+ op.drop_column("scans", "commit_hash")
+ op.drop_column("scans", "queued_by")
+ op.drop_column("scans", "queued_at")
+ op.drop_column("scans", "finished_by")
+ op.drop_column("scans", "status")
+ op.drop_column("scans", "pending_at")
+ # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.add_column("scans", sa.Column("pending_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=True))
+ op.add_column(
+ "scans",
+ sa.Column(
+ "status",
+ postgresql.ENUM("QUEUED", "PENDING", "FINISHED", "FAILED", name="status"),
+ autoincrement=False,
+ nullable=False,
+ ),
+ )
+ op.add_column("scans", sa.Column("finished_by", sa.VARCHAR(), autoincrement=False, nullable=True))
+ op.add_column("scans", sa.Column("queued_at", postgresql.TIMESTAMP(), autoincrement=False, nullable=True))
+ op.add_column("scans", sa.Column("queued_by", sa.VARCHAR(), autoincrement=False, nullable=False))
+ op.add_column("scans", sa.Column("commit_hash", sa.VARCHAR(), autoincrement=False, nullable=True))
+ op.add_column("scans", sa.Column("pending_by", sa.VARCHAR(), autoincrement=False, nullable=True))
+ op.add_column("scans", sa.Column("fail_reason", sa.VARCHAR(), autoincrement=False, nullable=True))
+ op.create_index(None, "scans", ["finished_at"])
+ op.create_table(
+ "download_urls",
+ sa.Column("id", sa.UUID(), autoincrement=False, nullable=False),
+ sa.Column("scan_id", sa.UUID(), autoincrement=False, nullable=False, index=True),
+ sa.Column("url", sa.VARCHAR(), autoincrement=False, nullable=False),
+ sa.ForeignKeyConstraint(["scan_id"], ["scans.scan_id"], name="download_urls_scan_id_fkey"),
+ sa.PrimaryKeyConstraint("id", name="download_urls_pkey"),
+ )
+ # ### end Alembic commands ###

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
  "requests>=2.30.0",
  "httpx>=0.24.1",
  "psycopg2-binary",
+ "aio-pika>=9.3.0",
 ]
 dynamic = ["version"]
 

src/mainframe/constants.py

@@ -31,7 +31,7 @@ class Mainframe(EnvConfig):
 
  dragonfly_github_token: str
 
- job_timeout: int = 60 * 2
+ amqp_connection_string: str = "amqp://127.0.0.1:5672/%2f"
 
 
 mainframe_settings = Mainframe() # pyright: ignore