Files
request-itv/request_itv/db.py
Ivan Vazhenin 1fe59867e7 .
2024-04-25 20:35:30 +03:00

167 lines
4.7 KiB
Python

from datetime import datetime
from typing import List, Optional
from sqlalchemy import create_engine, String, select, ForeignKey, Enum
from sqlalchemy.orm import Session, DeclarativeBase, Mapped, mapped_column, relationship
from config import Config
def tow(day: int, hour: int, minute: int):
return minute + hour * 60 + day * 60 * 24
class Base(DeclarativeBase):
pass
class User(Base):
__tablename__ = 'users'
id: Mapped[int] = mapped_column(primary_key=True)
username: Mapped[str]
passwd: Mapped[str]
bndname: Mapped[str]
newbnd: Mapped[bool]
active: Mapped[bool]
upstream: Mapped[bool]
profiles: Mapped[List['Profile']] = relationship(
back_populates='user', cascade='all, delete-orphan'
)
income_branches: Mapped[List['IncomeBranch']] = relationship(
back_populates='user', cascade='all, delete-orphan'
)
schedule: Mapped[List['Schedule']] = relationship(
back_populates='user', cascade='all, delete-orphan'
)
queue: Mapped[List['Queue']] = relationship(
back_populates='user', cascade='all, delete-orphan'
)
def __repr__(self) -> str:
return f'User(id={self.id!r}, username={self.username!r}, password={self.passwd!r}, newbnd={self.newbnd})'
def to_dict(self) -> dict:
return {
'id': self.id,
'username': self.username,
'bndname': self.bndname,
'newbnd': self.newbnd,
'active': self.active,
'upstream': self.upstream,
'profiles': [x.to_dict() for x in self.profiles],
'schedule': [x.to_dict() for x in self.schedule],
'queue': [x.to_dict() for x in self.queue],
'income_branches': [x.to_dict() for x in self.income_branches],
}
def is_active_now(self):
if not len(self.schedule):
return True
dt = datetime.now()
curr_tow = tow(dt.weekday(), dt.hour, dt.minute)
for x in self.schedule:
if (tow(x.day_start, x.hour_start, x.minute_start) <= curr_tow
<= tow(x.day_end, x.hour_end, x.minute_end)):
return True
return False
class Profile(Base):
__tablename__ = 'profiles'
id: Mapped[int] = mapped_column(primary_key=True)
user_id: Mapped[int] = mapped_column(ForeignKey('users.id'))
scheme: Mapped[str]
branch: Mapped[str] = mapped_column(String, nullable=True)
json: Mapped[str] = mapped_column(String, nullable=True)
no_files: Mapped[bool]
user: Mapped['User'] = relationship(back_populates='profiles')
def to_dict(self) -> dict:
return {
'id': self.id,
'scheme': self.scheme,
'branch': self.branch,
'json': self.json,
}
class IncomeBranch(Base):
__tablename__ = 'income_branches'
id: Mapped[int] = mapped_column(primary_key=True)
user_id: Mapped[int] = mapped_column(ForeignKey('users.id'))
scheme: Mapped[str]
branch: Mapped[str]
local_scheme: Mapped[str]
user: Mapped['User'] = relationship(back_populates='income_branches')
def to_dict(self) -> dict:
return {
'id': self.id,
'scheme': self.scheme,
'branch': self.branch,
'local_scheme': self.local_scheme,
}
class Schedule(Base):
__tablename__ = 'schedule'
id: Mapped[int] = mapped_column(primary_key=True)
user_id: Mapped[int] = mapped_column(ForeignKey('users.id'))
day_start: Mapped[int]
hour_start: Mapped[int]
minute_start: Mapped[int]
day_end: Mapped[int]
hour_end: Mapped[int]
minute_end: Mapped[int]
user: Mapped['User'] = relationship(back_populates='schedule')
def to_dict(self) -> dict:
return {
'id': self.id,
'day_start': self.day_start,
'hour_start': self.hour_start,
'minute_start': self.minute_start,
'day_end': self.day_end,
'hour_end': self.hour_end,
'minute_end': self.minute_end,
}
class Queue(Base):
__tablename__ = 'queue'
id: Mapped[int] = mapped_column(primary_key=True)
user_id: Mapped[int] = mapped_column(ForeignKey('users.id'))
commit_id: Mapped[str]
schema: Mapped[str]
user: Mapped['User'] = relationship(back_populates='queue')
def to_dict(self) -> dict:
return {
'id': self.id,
'commit_id': self.commit_id,
'schema': self.schema,
}
class Schemas(Base):
__tablename__ = 'schemas'
id: Mapped[int] = mapped_column(primary_key=True)
schema: Mapped[str]
schema_type: Mapped[str]
def connect_db():
return create_engine(f"postgresql+psycopg://{Config.pg_username}:{Config.pg_password}@{Config.pg_host}:{Config.pg_port}/{Config.pg_dbname}")