Skip to content

Commit

Permalink
Merge pull request #164 from vipyrsec/orm-dataclass-models
Browse files Browse the repository at this point in the history 
Convert SQLAlchemy models to dataclasses
  • Loading branch information
@Robin5605
Robin5605 committed Oct 31, 2023
2 parents ad3b0c5 + d483169 commit 59c6ec5
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 34 deletions.
62 changes: 35 additions & 27 deletions src/mainframe/models/orm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,16 @@
UniqueConstraint,
)
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from sqlalchemy.orm import (
DeclarativeBase,
Mapped,
MappedAsDataclass,
mapped_column,
relationship,
)


class Base(DeclarativeBase):
class Base(MappedAsDataclass, DeclarativeBase, kw_only=True):
pass


Expand Down Expand Up @@ -58,36 +64,38 @@ class Scan(Base):
scan_id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True),
primary_key=True,
server_default=FetchedValue(),
default=uuid.uuid4,
default_factory=uuid.uuid4,
init=False,
)

name: Mapped[str]
version: Mapped[str]
status: Mapped[Status]
name: Mapped[str] = mapped_column(default=None)
version: Mapped[str] = mapped_column(default=None)
status: Mapped[Status] = mapped_column(default=None)

score: Mapped[Optional[int]]
inspector_url: Mapped[Optional[str]]
rules: Mapped[list[Rule]] = relationship(secondary=package_rules)
download_urls: Mapped[list[DownloadURL]] = relationship()
score: Mapped[Optional[int]] = mapped_column(default=None)
inspector_url: Mapped[Optional[str]] = mapped_column(default=None)
rules: Mapped[list[Rule]] = relationship(secondary=package_rules, default_factory=list)
download_urls: Mapped[list[DownloadURL]] = relationship(default_factory=list)

queued_at: Mapped[Optional[datetime]] = mapped_column(
DateTime(timezone=True), server_default=FetchedValue(), default=lambda: datetime.now(timezone.utc)
DateTime(timezone=True),
server_default=FetchedValue(),
default_factory=lambda: datetime.now(timezone.utc),
)
queued_by: Mapped[str]
queued_by: Mapped[str] = mapped_column(default=None)

pending_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
pending_by: Mapped[Optional[str]]
pending_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), default=None)
pending_by: Mapped[Optional[str]] = mapped_column(default=None)

finished_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), index=True)
finished_by: Mapped[Optional[str]]
finished_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), default=None, index=True)
finished_by: Mapped[Optional[str]] = mapped_column(default=None)

reported_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
reported_by: Mapped[Optional[str]]
reported_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), default=None)
reported_by: Mapped[Optional[str]] = mapped_column(default=None)

fail_reason: Mapped[Optional[str]]
fail_reason: Mapped[Optional[str]] = mapped_column(default=None)

commit_hash: Mapped[Optional[str]]
commit_hash: Mapped[Optional[str]] = mapped_column(default=None)


Index(None, Scan.status, postgresql_where=or_(Scan.status == Status.QUEUED, Scan.status == Status.PENDING))
Expand All @@ -101,13 +109,13 @@ class DownloadURL(Base):
id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True),
primary_key=True,
server_default=FetchedValue(),
default=uuid.uuid4,
default_factory=uuid.uuid4,
init=False,
)

scan_id: Mapped[str] = mapped_column(ForeignKey("scans.scan_id"), index=True)
scan_id: Mapped[str] = mapped_column(ForeignKey("scans.scan_id"), index=True, init=False)

url: Mapped[str]
url: Mapped[str] = mapped_column()


class Rule(Base):
Expand All @@ -118,8 +126,8 @@ class Rule(Base):
id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True),
primary_key=True,
server_default=FetchedValue(),
default=uuid.uuid4,
default_factory=uuid.uuid4,
init=False,
)

name: Mapped[str] = mapped_column(unique=True)
6 changes: 3 additions & 3 deletions tests/test_data/first_safe_second_malicious.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

data = [
Scan(
scan_id="fce0366b-0bcf-4a29-a0a7-4d4bdf3c6f61",
name="a",
version="0.1.0",
status=Status.FINISHED,
Expand All @@ -19,11 +18,11 @@
reported_at=None,
reported_by=None,
rules=[],
download_urls=[],
commit_hash=None,
fail_reason=None,
),
Scan(
scan_id="13e22275-9053-4490-ab80-1800789d37da",
name="a",
version="0.2.0",
status=Status.FINISHED,
Expand All @@ -37,7 +36,8 @@
finished_by="remmy",
reported_at=None,
reported_by=None,
rules=[Rule(id="50bd2157-72ef-49c9-b75c-25e0b39cae5d", name="test rule 1")],
download_urls=[],
rules=[Rule(name="test rule 1")],
commit_hash="test commit hash",
fail_reason=None,
),
Expand Down
8 changes: 4 additions & 4 deletions tests/test_data/sample.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

data = [
Scan(
scan_id="fce0366b-0bcf-4a29-a0a7-4d4bdf3c6f61",
name="a",
version="0.1.0",
status=Status.FINISHED,
Expand All @@ -19,11 +18,11 @@
reported_at=None,
reported_by=None,
rules=[],
download_urls=[],
commit_hash=None,
fail_reason=None,
),
Scan(
scan_id="df157a3c-8994-467f-a494-9d63eaf96564",
name="b",
version="0.1.0",
status=Status.FINISHED,
Expand All @@ -38,11 +37,11 @@
reported_at=None,
reported_by=None,
rules=[],
download_urls=[],
commit_hash="test commit hash",
fail_reason=None,
),
Scan(
scan_id="04685768-e41d-49e4-9192-19b6d435226a",
name="a",
version="0.2.0",
status=Status.QUEUED,
Expand All @@ -57,11 +56,11 @@
reported_at=None,
reported_by=None,
rules=[],
download_urls=[],
commit_hash=None,
fail_reason=None,
),
Scan(
scan_id="bbb953ca-95af-4d57-a7a5-0b656f652695",
name="b",
version="0.2.0",
status=Status.PENDING,
Expand All @@ -76,6 +75,7 @@
reported_at=None,
reported_by=None,
rules=[],
download_urls=[],
commit_hash=None,
fail_reason=None,
),
Expand Down

0 comments on commit 59c6ec5

Please sign in to comment.