diff --git a/requirements.txt b/requirements.txt index 27da0663..eeefe2b2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,7 @@ redis==5.0.8 requests==2.32.2 rq==1.16.2 sniffio==1.3.1 -sqlmodel==0.0.11 +sqlmodel==0.0.22 starlette==0.38.6 typed-config==2.0.3 types-requests==2.32.0.20240914 diff --git a/src/models/jobagent.py b/src/models/jobagent.py index a8041ea3..1d94cf85 100644 --- a/src/models/jobagent.py +++ b/src/models/jobagent.py @@ -1,4 +1,3 @@ -from sqlalchemy import Column, ForeignKey from sqlmodel import SQLModel, Field, Relationship from typing import Union, TYPE_CHECKING @@ -13,9 +12,7 @@ class JobAgent(SQLModel, table=True): id: Union[int, None] = Field(default=None, primary_key=True) task_in_progress: int - job_id: int = Field( - sa_column=Column(ForeignKey("job.internal_id", ondelete="CASCADE")) - ) + job_id: int = Field(foreign_key="job.internal_id", ondelete="CASCADE") job: "Job" = Relationship(back_populates="agents") agent_id: int = Field(foreign_key="agentgroup.id") diff --git a/src/models/match.py b/src/models/match.py index 7e089688..0191f655 100644 --- a/src/models/match.py +++ b/src/models/match.py @@ -1,4 +1,3 @@ -from sqlalchemy import ForeignKey from sqlmodel import SQLModel, Field, ARRAY, String, Column, JSON, Relationship from typing import List, Union, Dict, Any @@ -16,7 +15,5 @@ class Match(SQLModel, table=True): # A list of yara rules matched to this file matches: List[str] = Field(sa_column=Column(ARRAY(String))) - job_id: int = Field( - sa_column=Column(ForeignKey("job.internal_id", ondelete="CASCADE")) - ) + job_id: int = Field(foreign_key="job.internal_id", ondelete="CASCADE") job: Job = Relationship(back_populates="matches")