Add Distribution model

alembic/versions/a62a93704798_add_distributions.py

@@ -0,0 +1,23 @@
+"""add distributions
+
+Revision ID: a62a93704798
+Revises: 587c186d91ee
+Create Date: 2024-08-11 08:12:42.354151
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "a62a93704798"
+down_revision = "587c186d91ee"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.alter_column("scans", "files", new_column_name="distributions")
+
+
+def downgrade() -> None:
+    op.alter_column("scans", "distributions", new_column_name="files")

src/mainframe/endpoints/package.py

@@ -78,7 +78,7 @@ def submit_results(
         scan.score = result.score
         scan.finished_by = auth.subject
         scan.commit_hash = result.commit
-        scan.files = result.files
+        scan.distributions = result.distributions
 
         # These are the rules that already have an entry in the database
         rules = session.scalars(select(Rule).where(Rule.name.in_(result.rules_matched))).all()

src/mainframe/models/orm.py

@@ -28,7 +28,7 @@
 )
 
 from mainframe.models import Pydantic
-from mainframe.models.schemas import Files
+from mainframe.models.schemas import Distributions
 
 
 class Base(MappedAsDataclass, DeclarativeBase, kw_only=True):
@@ -102,7 +102,7 @@ class Scan(Base):
 
     commit_hash: Mapped[Optional[str]] = mapped_column(default=None)
 
-    files: Mapped[Optional[Files]] = mapped_column(Pydantic(Files), default=None)
+    distributions: Mapped[Optional[Distributions]] = mapped_column(Pydantic(Distributions), default=None)
 
 
 Index(None, Scan.status, postgresql_where=or_(Scan.status == Status.QUEUED, Scan.status == Status.PENDING))

src/mainframe/models/schemas.py

@@ -48,7 +48,15 @@ class File(BaseModel):
     matches: list[RuleMatch]
 
 
-Files = RootModel[list[File]]
+Files = list[File]
+
+
+class Distribution(BaseModel):
+    download_url: str
+    files: Files
+
+
+Distributions = RootModel[list[Distribution]]
 
 
 class ServerMetadata(BaseModel):
@@ -88,7 +96,7 @@ class Package(BaseModel):
 
     commit_hash: Optional[str]
 
-    files: Optional[Files]
+    distributions: Optional[Distributions]
 
     @classmethod
     def from_db(cls, scan: Scan):
@@ -110,7 +118,7 @@ def from_db(cls, scan: Scan):
             finished_at=scan.finished_at,
             finished_by=scan.finished_by,
             commit_hash=scan.commit_hash,
-            files=scan.files,
+            distributions=scan.distributions,
         )
 
     @field_serializer(
@@ -179,7 +187,7 @@ class PackageScanResult(PackageSpecifier):
     score: int = 0
     inspector_url: Optional[str] = None
     rules_matched: list[str] = []
-    files: Optional[Files] = None
+    distributions: Optional[Distributions] = None
 
 
 class PackageScanResultFail(PackageSpecifier):

tests/test_package.py

@@ -20,6 +20,8 @@
 from mainframe.json_web_token import AuthenticationData
 from mainframe.models.orm import Scan, Status
 from mainframe.models.schemas import (
+    Distribution,
+    Distributions,
     File,
     Files,
     Match,
@@ -95,21 +97,21 @@ def test_package_lookup_files(db_session: Session):
     rule = RuleMatch(identifier="rule1", patterns=[pattern], metadata={"author": "remmy", "score": 5})
     file = File(path="dist1/a/b.py", matches=[rule])
     files = Files([file])
+    distros = Distributions([Distribution(download_url="http://example.com", files=files)])
     scan = Scan(
         name="abc",
         version="1.0.0",
         status=Status.FINISHED,
         queued_by="remmy",
-        files=files,
+        distributions=distros,
     )
 
     with db_session.begin():
         db_session.add(scan)
-        db_session.commit()
 
     package = lookup_package_info(db_session, name="abc", version="1.0.0")[0]
 
-    assert package.files == files
+    assert package.distributions == distros
 
 
 def test_handle_success(db_session: Session, test_data: list[Scan], auth: AuthenticationData, rules_state: Rules):
@@ -126,6 +128,7 @@ def test_handle_success(db_session: Session, test_data: list[Scan], auth: Authen
         rule = RuleMatch(identifier="rule1", patterns=[pattern], metadata={"author": "remmy", "score": 5})
         file = File(path="dist1/a/b.py", matches=[rule])
         files = Files([file])
+        distros = Distributions([Distribution(download_url="http://example.com", files=files)])
 
         body = PackageScanResult(
             name=job.name,
@@ -134,7 +137,7 @@ def test_handle_success(db_session: Session, test_data: list[Scan], auth: Authen
             score=2,
             inspector_url="test inspector url",
             rules_matched=["a", "b", "c"],
-            files=files,
+            distributions=distros,
         )
         submit_results(body, db_session, auth)
 
@@ -147,7 +150,7 @@ def test_handle_success(db_session: Session, test_data: list[Scan], auth: Authen
         assert record.score == 2
         assert record.inspector_url == "test inspector url"
         assert {rule.name for rule in record.rules} == {"a", "b", "c"}
-        assert record.files == files
+        assert record.distributions == distros
     else:
         assert all(scan.status != Status.QUEUED for scan in test_data)