45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
|
import os
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from typing import Iterable
|
||
|
|
|
||
|
|
from sqlalchemy import Column, Integer, String, create_engine, select
|
||
|
|
from sqlalchemy.exc import SQLAlchemyError
|
||
|
|
from sqlalchemy.orm import declarative_base, sessionmaker
|
||
|
|
|
||
|
|
DATABASE_URL = os.getenv("DATABASE_URL", "mysql+pymysql://ec2_mt5:8FmzXj4xcz3AiH2R@163.123.183.106:3306/ip_ops")
|
||
|
|
|
||
|
|
engine = create_engine(DATABASE_URL, pool_pre_ping=True)
|
||
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
||
|
|
Base = declarative_base()
|
||
|
|
|
||
|
|
|
||
|
|
class IPOperation(Base):
|
||
|
|
__tablename__ = "ip_operations"
|
||
|
|
|
||
|
|
id = Column(Integer, primary_key=True, autoincrement=True)
|
||
|
|
ip_address = Column(String(64), unique=True, nullable=False, index=True)
|
||
|
|
note = Column(String(255), nullable=True)
|
||
|
|
|
||
|
|
|
||
|
|
def init_db() -> None:
|
||
|
|
Base.metadata.create_all(bind=engine)
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def db_session():
|
||
|
|
session = SessionLocal()
|
||
|
|
try:
|
||
|
|
yield session
|
||
|
|
session.commit()
|
||
|
|
except SQLAlchemyError:
|
||
|
|
session.rollback()
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
session.close()
|
||
|
|
|
||
|
|
|
||
|
|
def load_disallowed_ips() -> set[str]:
|
||
|
|
with db_session() as session:
|
||
|
|
rows: Iterable[IPOperation] = session.scalars(select(IPOperation.ip_address))
|
||
|
|
return {row for row in rows}
|