refactor: moving more models to mapped_column

This commit is contained in:
Iain Learmonth 2024-11-10 15:13:29 +00:00
parent ea020d6edd
commit 75b2c1adf0
9 changed files with 272 additions and 94 deletions

View file

@ -2,10 +2,10 @@ from __future__ import annotations
import json
from datetime import datetime, timedelta
from typing import Optional, List, Union, Any, Dict
from typing import Optional, List, Union, Any, Dict, TypedDict, Literal
import tldextract
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.orm import Mapped, mapped_column, relationship
from tldextract import extract
from werkzeug.datastructures import FileStorage
@ -13,7 +13,7 @@ from app.brm.brn import BRN
from app.brm.utils import thumbnail_uploaded_image, create_data_uri, normalize_color
from app.extensions import db
from app.models import AbstractConfiguration, AbstractResource, Deprecation
from app.models.base import Pool
from app.models.base import Pool, Group
from app.models.onions import Onion
country_origin = db.Table(
@ -25,17 +25,25 @@ country_origin = db.Table(
)
class Origin(AbstractConfiguration):
group_id = mapped_column(db.Integer, db.ForeignKey("group.id"), nullable=False)
domain_name = mapped_column(db.String(255), unique=True, nullable=False)
auto_rotation = mapped_column(db.Boolean, nullable=False)
smart = mapped_column(db.Boolean(), nullable=False)
assets = mapped_column(db.Boolean(), nullable=False)
risk_level_override = mapped_column(db.Integer(), nullable=True)
class OriginDict(TypedDict):
Id: int
Description: str
DomainName: str
RiskLevel: Dict[str, int]
RiskLevelOverride: Optional[int]
group = db.relationship("Group", back_populates="origins")
proxies = db.relationship("Proxy", back_populates="origin")
countries = db.relationship("Country", secondary=country_origin, back_populates='origins')
class Origin(AbstractConfiguration):
group_id: Mapped[int] = mapped_column(db.Integer, db.ForeignKey("group.id"))
domain_name: Mapped[str] = mapped_column(db.String(255), unique=True)
auto_rotation: Mapped[bool]
smart: Mapped[bool]
assets: Mapped[bool]
risk_level_override: Mapped[Optional[int]]
group: Mapped[Group] = relationship("Group", back_populates="origins")
proxies: Mapped[List[Proxy]] = relationship("Proxy", back_populates="origin")
countries: Mapped[List[Country]] = relationship("Country", secondary=country_origin, back_populates='origins')
@property
def brn(self) -> BRN:
@ -100,7 +108,7 @@ class Origin(AbstractConfiguration):
max(1, min(10, frequency_factor * recency_factor))) + country.risk_level
return risk_levels
def to_dict(self):
def to_dict(self) -> OriginDict:
return {
"Id": self.id,
"Description": self.description,
@ -250,6 +258,16 @@ class StaticOrigin(AbstractConfiguration):
self.updated = datetime.utcnow()
ResourceStatus = Union[Literal["active"], Literal["pending"], Literal["expiring"], Literal["destroyed"]]
class ProxyDict(TypedDict):
Id: int
OriginDomain: str
MirrorDomain: Optional[str]
Status: ResourceStatus
class Proxy(AbstractResource):
origin_id: Mapped[int] = mapped_column(db.Integer, db.ForeignKey("origin.id"), nullable=False)
pool_id: Mapped[Optional[int]] = mapped_column(db.Integer, db.ForeignKey("pool.id"))
@ -259,8 +277,8 @@ class Proxy(AbstractResource):
terraform_updated: Mapped[Optional[datetime]] = mapped_column(db.DateTime(), nullable=True)
url: Mapped[Optional[str]] = mapped_column(db.String(255), nullable=True)
origin: Mapped[Origin] = db.relationship("Origin", back_populates="proxies")
pool: Mapped[Pool] = db.relationship("Pool", back_populates="proxies")
origin: Mapped[Origin] = relationship("Origin", back_populates="proxies")
pool: Mapped[Pool] = relationship("Pool", back_populates="proxies")
@property
def brn(self) -> BRN:
@ -278,8 +296,8 @@ class Proxy(AbstractResource):
"origin_id", "provider", "psg", "slug", "terraform_updated", "url"
]
def to_dict(self):
status = "active"
def to_dict(self) -> ProxyDict:
status: ResourceStatus = "active"
if self.url is None:
status = "pending"
if self.deprecated is not None: