mirror of
https://github.com/EstrellaXD/Auto_Bangumi.git
synced 2026-07-09 05:28:04 +08:00
Merge branch '3.1-dev' into self-rss
# Conflicts: # .gitignore # backend/src/module/database/bangumi.py # backend/src/module/database/orm/update.py # backend/src/module/models/__init__.py # backend/src/module/rss/poller.py
This commit is contained in:
estrella
@@ -1,138 +1,102 @@
|
||||
import logging
|
||||
|
||||
from module.database.orm import Connector
|
||||
from module.models import BangumiData
|
||||
from module.conf import DATA_PATH
|
||||
from sqlmodel import Session, select, delete, or_
|
||||
from sqlalchemy.sql import func
|
||||
from typing import Optional
|
||||
|
||||
from .engine import engine
|
||||
from module.models import Bangumi
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BangumiDatabase(Connector):
|
||||
def __init__(self, database: str = DATA_PATH):
|
||||
super().__init__(
|
||||
table_name="bangumi",
|
||||
data=self.__data_to_db(BangumiData()),
|
||||
database=database,
|
||||
)
|
||||
class BangumiDatabase(Session):
|
||||
def __init__(self, _engine=engine):
|
||||
super().__init__(_engine)
|
||||
|
||||
def update_table(self):
|
||||
self.update.table()
|
||||
|
||||
@staticmethod
|
||||
def __data_to_db(data: BangumiData) -> dict:
|
||||
db_data = data.dict()
|
||||
for key, value in db_data.items():
|
||||
if isinstance(value, bool):
|
||||
db_data[key] = int(value)
|
||||
elif isinstance(value, list):
|
||||
db_data[key] = ",".join(value)
|
||||
return db_data
|
||||
|
||||
@staticmethod
|
||||
def __db_to_data(db_data: dict) -> BangumiData:
|
||||
for key, item in db_data.items():
|
||||
if isinstance(item, int):
|
||||
if key not in ["id", "offset", "season", "year"]:
|
||||
db_data[key] = bool(item)
|
||||
elif key in ["filter", "rss_link"]:
|
||||
db_data[key] = item.split(",")
|
||||
return BangumiData(**db_data)
|
||||
|
||||
def insert_one(self, data: BangumiData):
|
||||
db_data = self.__data_to_db(data)
|
||||
self.insert.one(db_data)
|
||||
def insert_one(self, data: Bangumi):
|
||||
self.add(data)
|
||||
self.commit()
|
||||
logger.debug(f"[Database] Insert {data.official_title} into database.")
|
||||
# if self.__check_exist(data):
|
||||
# self.update_one(data)
|
||||
# else:
|
||||
# db_data = self.__data_to_db(data)
|
||||
# db_data["id"] = self.gen_id()
|
||||
# self._insert(db_data=db_data, table_name=self.__table_name)
|
||||
# logger.debug(f"[Database] Insert {data.official_title} into database.")
|
||||
|
||||
def insert_list(self, data: list[BangumiData]):
|
||||
data_list = [self.__data_to_db(x) for x in data]
|
||||
self.insert.many(data_list)
|
||||
# _id = self.gen_id()
|
||||
# for i, item in enumerate(data):
|
||||
# item.id = _id + i
|
||||
# data_list = [self.__data_to_db(x) for x in data]
|
||||
# self._insert_list(data_list=data_list, table_name=self.__table_name)
|
||||
def insert_list(self, data: list[Bangumi]):
|
||||
self.add_all(data)
|
||||
logger.debug(f"[Database] Insert {len(data)} bangumi into database.")
|
||||
|
||||
def update_one(self, data: BangumiData) -> bool:
|
||||
db_data = self.__data_to_db(data)
|
||||
return self.update.one(db_data)
|
||||
def update_one(self, data: Bangumi) -> bool:
|
||||
db_data = self.get(Bangumi, data.id)
|
||||
if not db_data:
|
||||
return False
|
||||
bangumi_data = data.dict(exclude_unset=True)
|
||||
for key, value in bangumi_data.items():
|
||||
setattr(db_data, key, value)
|
||||
self.add(db_data)
|
||||
self.commit()
|
||||
self.refresh(db_data)
|
||||
logger.debug(f"[Database] Update {data.official_title}")
|
||||
return True
|
||||
|
||||
def update_list(self, data: list[BangumiData]):
|
||||
data_list = [self.__data_to_db(x) for x in data]
|
||||
self.update.many(data_list)
|
||||
def update_list(self, datas: list[Bangumi]):
|
||||
for data in datas:
|
||||
self.update_one(data)
|
||||
|
||||
def update_rss(self, title_raw, rss_set: str):
|
||||
# Update rss and added
|
||||
location = {"title_raw": title_raw}
|
||||
set_value = {"rss_link": rss_set, "added": 0}
|
||||
self.update.value(location, set_value)
|
||||
# self._cursor.execute(
|
||||
# """
|
||||
# UPDATE bangumi
|
||||
# SET rss_link = :rss_link, added = 0
|
||||
# WHERE title_raw = :title_raw
|
||||
# """,
|
||||
# {"rss_link": rss_set, "title_raw": title_raw},
|
||||
# )
|
||||
# self._conn.commit()
|
||||
statement = select(Bangumi).where(Bangumi.title_raw == title_raw)
|
||||
bangumi = self.exec(statement).first()
|
||||
bangumi.rss_link = rss_set
|
||||
bangumi.added = False
|
||||
self.add(bangumi)
|
||||
self.commit()
|
||||
self.refresh(bangumi)
|
||||
logger.debug(f"[Database] Update {title_raw} rss_link to {rss_set}.")
|
||||
|
||||
def update_poster(self, title_raw, poster_link: str):
|
||||
location = {"title_raw": title_raw}
|
||||
set_value = {"poster_link": poster_link}
|
||||
self.update.value(location, set_value)
|
||||
statement = select(Bangumi).where(Bangumi.title_raw == title_raw)
|
||||
bangumi = self.exec(statement).first()
|
||||
bangumi.poster_link = poster_link
|
||||
self.add(bangumi)
|
||||
self.commit()
|
||||
self.refresh(bangumi)
|
||||
logger.debug(f"[Database] Update {title_raw} poster_link to {poster_link}.")
|
||||
|
||||
def delete_one(self, _id: int):
|
||||
self.delete.one(_id)
|
||||
statement = select(Bangumi).where(Bangumi.id == _id)
|
||||
bangumi = self.exec(statement).first()
|
||||
self.delete(bangumi)
|
||||
self.commit()
|
||||
logger.debug(f"[Database] Delete bangumi id: {_id}.")
|
||||
|
||||
def delete_all(self):
|
||||
self.delete.all()
|
||||
statement = delete(Bangumi)
|
||||
self.exec(statement)
|
||||
self.commit()
|
||||
|
||||
def search_all(self) -> list[BangumiData]:
|
||||
all_data = self.select.all()
|
||||
return [self.__db_to_data(x) for x in all_data]
|
||||
def search_all(self) -> list[Bangumi]:
|
||||
statement = select(Bangumi)
|
||||
return self.exec(statement).all()
|
||||
|
||||
def search_id(self, _id: int) -> BangumiData | None:
|
||||
dict_data = self.select.one(conditions={"id": _id})
|
||||
if dict_data is None:
|
||||
def search_id(self, _id: int) -> Optional[Bangumi]:
|
||||
statement = select(Bangumi).where(Bangumi.id == _id)
|
||||
bangumi = self.exec(statement).first()
|
||||
if bangumi is None:
|
||||
logger.warning(f"[Database] Cannot find bangumi id: {_id}.")
|
||||
return None
|
||||
logger.debug(f"[Database] Find bangumi id: {_id}.")
|
||||
return self.__db_to_data(dict_data)
|
||||
|
||||
# def search_official_title(self, official_title: str) -> BangumiData | None:
|
||||
# dict_data = self._search_data(
|
||||
# table_name=self.__table_name, condition={"official_title": official_title}
|
||||
# )
|
||||
# if dict_data is None:
|
||||
# return None
|
||||
# return self.__db_to_data(dict_data)
|
||||
else:
|
||||
logger.debug(f"[Database] Find bangumi id: {_id}.")
|
||||
return self.exec(statement).first()
|
||||
|
||||
def match_poster(self, bangumi_name: str) -> str:
|
||||
condition = {"official_title": bangumi_name}
|
||||
keys = ["poster_link"]
|
||||
data = self.select.one(
|
||||
keys=keys,
|
||||
conditions=condition,
|
||||
combine_operator="INSTR",
|
||||
)
|
||||
if not data:
|
||||
# Use like to match
|
||||
statement = select(Bangumi).where(func.instr(bangumi_name, Bangumi.title_raw) > 0)
|
||||
data = self.exec(statement).first()
|
||||
if data:
|
||||
return data.poster_link
|
||||
else:
|
||||
return ""
|
||||
return data.get("poster_link")
|
||||
|
||||
def match_list(self, torrent_list: list, rss_link: str) -> list:
|
||||
# Match title_raw in database
|
||||
keys = ["title_raw", "rss_link", "poster_link"]
|
||||
match_datas = self.select.column(keys)
|
||||
match_datas = self.search_all()
|
||||
if not match_datas:
|
||||
return torrent_list
|
||||
# Match title
|
||||
@@ -140,50 +104,43 @@ class BangumiDatabase(Connector):
|
||||
while i < len(torrent_list):
|
||||
torrent = torrent_list[i]
|
||||
for match_data in match_datas:
|
||||
if match_data.get("title_raw") in torrent.name:
|
||||
if rss_link not in match_data.get("rss_link"):
|
||||
match_data["rss_link"] += f",{rss_link}"
|
||||
self.update_rss(
|
||||
match_data.get("title_raw"), match_data.get("rss_link")
|
||||
)
|
||||
if not match_data.get("poster_link"):
|
||||
self.update_poster(
|
||||
match_data.get("title_raw"), torrent.poster_link
|
||||
)
|
||||
if match_data.title_raw in torrent.name:
|
||||
if rss_link not in match_data.rss_link:
|
||||
match_data.rss_link += f",{rss_link}"
|
||||
self.update_rss(match_data.title_raw, match_data.rss_link)
|
||||
if not match_data.poster_link:
|
||||
self.update_poster(match_data.title_raw, torrent.poster_link)
|
||||
torrent_list.pop(i)
|
||||
break
|
||||
else:
|
||||
i += 1
|
||||
return torrent_list
|
||||
|
||||
def not_complete(self) -> list[BangumiData]:
|
||||
def not_complete(self) -> list[Bangumi]:
|
||||
# Find eps_complete = False
|
||||
condition = {"eps_collect": 0}
|
||||
dict_data = self.select.many(
|
||||
conditions=condition,
|
||||
condition = select(Bangumi).where(Bangumi.eps_collect == 0)
|
||||
datas = self.exec(condition).all()
|
||||
return datas
|
||||
|
||||
def not_added(self) -> list[Bangumi]:
|
||||
conditions = select(Bangumi).where(
|
||||
or_(
|
||||
Bangumi.added == 0, Bangumi.rule_name is None, Bangumi.save_path is None
|
||||
)
|
||||
)
|
||||
return [self.__db_to_data(x) for x in dict_data]
|
||||
datas = self.exec(conditions).all()
|
||||
return datas
|
||||
|
||||
def not_added(self) -> list[BangumiData]:
|
||||
conditions = {"added": 0, "rule_name": None, "save_path": None}
|
||||
dict_data = self.select.many(conditions=conditions, combine_operator="OR")
|
||||
return [self.__db_to_data(x) for x in dict_data]
|
||||
|
||||
def get_rss(self, rss_link: str) -> list[BangumiData]:
|
||||
conditions = {"rss_link": rss_link}
|
||||
dict_data = self.select.many(conditions=conditions, combine_operator="INSTR")
|
||||
return [self.__db_to_data(x) for x in dict_data]
|
||||
|
||||
def match_torrent(self, torrent_name: str, rss_link: str) -> BangumiData | None:
|
||||
conditions = {"title_raw": torrent_name, "rss_link": rss_link}
|
||||
dict_data = self.select.one(conditions=conditions, combine_operator="INSTR")
|
||||
if not dict_data:
|
||||
return None
|
||||
return self.__db_to_data(dict_data)
|
||||
def disable_rule(self, _id: int):
|
||||
statement = select(Bangumi).where(Bangumi.id == _id)
|
||||
bangumi = self.exec(statement).first()
|
||||
bangumi.deleted = True
|
||||
self.add(bangumi)
|
||||
self.commit()
|
||||
self.refresh(bangumi)
|
||||
logger.debug(f"[Database] Disable rule {bangumi.title_raw}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
with BangumiDatabase() as db:
|
||||
db.match_torrent(
|
||||
"魔法科高校の劣等生 来訪者編", "https://bangumi.moe/rss/5f6b3e3e4e8c4b0001b2e3a3"
|
||||
)
|
||||
print(db.not_complete())
|
||||
|
||||
@@ -1,165 +0,0 @@
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
|
||||
from module.conf import DATA_PATH
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DataConnector:
|
||||
def __init__(self):
|
||||
# Create folder if not exists
|
||||
if not os.path.exists(os.path.dirname(DATA_PATH)):
|
||||
os.makedirs(os.path.dirname(DATA_PATH))
|
||||
self._conn = sqlite3.connect(DATA_PATH)
|
||||
self._cursor = self._conn.cursor()
|
||||
|
||||
def _update_table(self, table_name: str, db_data: dict):
|
||||
columns = ", ".join(
|
||||
[
|
||||
f"{key} {self.__python_to_sqlite_type(value)}"
|
||||
for key, value in db_data.items()
|
||||
]
|
||||
)
|
||||
create_table_sql = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns});"
|
||||
self._cursor.execute(create_table_sql)
|
||||
self._cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
existing_columns = {
|
||||
column_info[1]: column_info for column_info in self._cursor.fetchall()
|
||||
}
|
||||
for key, value in db_data.items():
|
||||
if key not in existing_columns:
|
||||
insert_column = self.__python_to_sqlite_type(value)
|
||||
if value is None:
|
||||
value = "NULL"
|
||||
add_column_sql = f"ALTER TABLE {table_name} ADD COLUMN {key} {insert_column} DEFAULT {value};"
|
||||
self._cursor.execute(add_column_sql)
|
||||
self._conn.commit()
|
||||
logger.debug(f"Create / Update table {table_name}.")
|
||||
|
||||
def _insert(self, table_name: str, db_data: dict):
|
||||
columns = ", ".join(db_data.keys())
|
||||
values = ", ".join([f":{key}" for key in db_data.keys()])
|
||||
self._cursor.execute(
|
||||
f"INSERT INTO {table_name} ({columns}) VALUES ({values})", db_data
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def _insert_list(self, table_name: str, data_list: list[dict]):
|
||||
columns = ", ".join(data_list[0].keys())
|
||||
values = ", ".join([f":{key}" for key in data_list[0].keys()])
|
||||
self._cursor.executemany(
|
||||
f"INSERT INTO {table_name} ({columns}) VALUES ({values})", data_list
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def _select(self, keys: list[str], table_name: str, condition: str = None) -> dict:
|
||||
if condition is None:
|
||||
self._cursor.execute(f"SELECT {', '.join(keys)} FROM {table_name}")
|
||||
else:
|
||||
self._cursor.execute(
|
||||
f"SELECT {', '.join(keys)} FROM {table_name} WHERE {condition}"
|
||||
)
|
||||
return dict(zip(keys, self._cursor.fetchone()))
|
||||
|
||||
def _update(self, table_name: str, db_data: dict):
|
||||
_id = db_data.get("id")
|
||||
if _id is None:
|
||||
raise ValueError("No _id in db_data.")
|
||||
set_sql = ", ".join([f"{key} = :{key}" for key in db_data.keys()])
|
||||
self._cursor.execute(
|
||||
f"UPDATE {table_name} SET {set_sql} WHERE id = {_id}", db_data
|
||||
)
|
||||
self._conn.commit()
|
||||
return self._cursor.rowcount == 1
|
||||
|
||||
def _update_list(self, table_name: str, data_list: list[dict]):
|
||||
if len(data_list) == 0:
|
||||
return
|
||||
set_sql = ", ".join(
|
||||
[f"{key} = :{key}" for key in data_list[0].keys() if key != "id"]
|
||||
)
|
||||
self._cursor.executemany(
|
||||
f"UPDATE {table_name} SET {set_sql} WHERE id = :id", data_list
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def _update_section(self, table_name: str, location: dict, update_dict: dict):
|
||||
set_sql = ", ".join([f"{key} = :{key}" for key in update_dict.keys()])
|
||||
sql_loc = f"{location['key']} = {location['value']}"
|
||||
self._cursor.execute(
|
||||
f"UPDATE {table_name} SET {set_sql} WHERE {sql_loc}", update_dict
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def _delete_all(self, table_name: str):
|
||||
self._cursor.execute(f"DELETE FROM {table_name}")
|
||||
self._conn.commit()
|
||||
|
||||
def _delete(self, table_name: str, condition: dict):
|
||||
condition_sql = " AND ".join([f"{key} = :{key}" for key in condition.keys()])
|
||||
self._cursor.execute(f"DELETE FROM {table_name} WHERE {condition_sql}", condition)
|
||||
self._conn.commit()
|
||||
|
||||
def _search(self, table_name: str, keys: list[str] | None = None, condition: dict = None):
|
||||
if keys is None:
|
||||
select_sql = "*"
|
||||
else:
|
||||
select_sql = ", ".join(keys)
|
||||
if condition is None:
|
||||
self._cursor.execute(f"SELECT {select_sql} FROM {table_name}")
|
||||
else:
|
||||
custom_condition = condition.pop("_custom_condition", None)
|
||||
condition_sql = " AND ".join([f"{key} = :{key}" for key in condition.keys()]) + (
|
||||
f" AND {custom_condition}" if custom_condition else ""
|
||||
)
|
||||
self._cursor.execute(
|
||||
f"SELECT {select_sql} FROM {table_name} WHERE {condition_sql}", condition
|
||||
)
|
||||
|
||||
def _search_data(self, table_name: str, keys: list[str] | None = None, condition: dict = None) -> dict:
|
||||
if keys is None:
|
||||
keys = self.__get_table_columns(table_name)
|
||||
self._search(table_name, keys, condition)
|
||||
return dict(zip(keys, self._cursor.fetchone()))
|
||||
|
||||
def _search_datas(self, table_name: str, keys: list[str] | None = None, condition: dict = None) -> list[dict]:
|
||||
if keys is None:
|
||||
keys = self.__get_table_columns(table_name)
|
||||
self._search(table_name, keys, condition)
|
||||
return [dict(zip(keys, row)) for row in self._cursor.fetchall()]
|
||||
|
||||
def _table_exists(self, table_name: str) -> bool:
|
||||
self._cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name=?;",
|
||||
(table_name,),
|
||||
)
|
||||
return len(self._cursor.fetchall()) == 1
|
||||
|
||||
def __get_table_columns(self, table_name: str) -> list[str]:
|
||||
self._cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
return [column_info[1] for column_info in self._cursor.fetchall()]
|
||||
|
||||
@staticmethod
|
||||
def __python_to_sqlite_type(value) -> str:
|
||||
if isinstance(value, int):
|
||||
return "INTEGER NOT NULL"
|
||||
elif isinstance(value, float):
|
||||
return "REAL NOT NULL"
|
||||
elif isinstance(value, str):
|
||||
return "TEXT NOT NULL"
|
||||
elif isinstance(value, bool):
|
||||
return "INTEGER NOT NULL"
|
||||
elif isinstance(value, list):
|
||||
return "TEXT NOT NULL"
|
||||
elif value is None:
|
||||
return "TEXT"
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {type(value)}")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self._conn.close()
|
||||
7
backend/src/module/database/engine.py
Normal file
7
backend/src/module/database/engine.py
Normal file
@@ -0,0 +1,7 @@
|
||||
from sqlmodel import create_engine, Session
|
||||
from module.conf import DATA_PATH
|
||||
|
||||
|
||||
engine = create_engine(DATA_PATH)
|
||||
|
||||
db_session = Session(engine)
|
||||
@@ -1 +0,0 @@
|
||||
from .connector import Connector
|
||||
@@ -1,62 +0,0 @@
|
||||
import sqlite3
|
||||
|
||||
from .delete import Delete
|
||||
from .insert import Insert
|
||||
from .select import Select
|
||||
from .update import Update
|
||||
|
||||
from module.conf import DATA_PATH
|
||||
|
||||
|
||||
class Connector:
|
||||
def __init__(self, table_name: str, data: dict, database: str = DATA_PATH):
|
||||
self._conn = sqlite3.connect(database)
|
||||
self._cursor = self._conn.cursor()
|
||||
self.update = Update(self, table_name, data)
|
||||
self.insert = Insert(self, table_name, data)
|
||||
self.select = Select(self, table_name, data)
|
||||
self.delete = Delete(self, table_name, data)
|
||||
self._columns = self.__get_columns(table_name)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self._conn.close()
|
||||
|
||||
def __get_columns(self, table_name: str) -> list[str]:
|
||||
self._cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
return [x[1] for x in self._cursor.fetchall()]
|
||||
|
||||
def execute(self, sql: str, params: tuple = None):
|
||||
if params is None:
|
||||
self._cursor.execute(sql)
|
||||
else:
|
||||
self._cursor.execute(sql, params)
|
||||
self._conn.commit()
|
||||
|
||||
def executemany(self, sql: str, params: list[tuple]):
|
||||
self._cursor.executemany(sql, params)
|
||||
self._conn.commit()
|
||||
|
||||
def fetchall(self, keys: str = None) -> list[dict]:
|
||||
datas = self._cursor.fetchall()
|
||||
if keys:
|
||||
return [dict(zip(keys, data)) for data in datas]
|
||||
return [dict(zip(self._columns, data)) for data in datas]
|
||||
|
||||
def fetchone(self, keys: list[str] = None) -> dict:
|
||||
data = self._cursor.fetchone()
|
||||
if data:
|
||||
if keys:
|
||||
return dict(zip(keys, data))
|
||||
return dict(zip(self._columns, data))
|
||||
|
||||
def fetchmany(self, keys: list[str], size: int) -> list[dict]:
|
||||
datas = self._cursor.fetchmany(size)
|
||||
if keys:
|
||||
return [dict(zip(keys, data)) for data in datas]
|
||||
return [dict(zip(self._columns, data)) for data in datas]
|
||||
|
||||
def fetch(self):
|
||||
return self._cursor.fetchall()
|
||||
@@ -1,23 +0,0 @@
|
||||
class Delete:
|
||||
def __init__(self, connector, table_name: str, data: dict):
|
||||
self._connector = connector
|
||||
self._table_name = table_name
|
||||
self._data = data
|
||||
|
||||
def one(self, _id: int) -> bool:
|
||||
self._connector.execute(
|
||||
f"""
|
||||
DELETE FROM {self._table_name}
|
||||
WHERE id = :id
|
||||
""",
|
||||
{"id": _id},
|
||||
)
|
||||
return True
|
||||
|
||||
def all(self):
|
||||
self._connector.execute(
|
||||
f"""
|
||||
DELETE FROM {self._table_name}
|
||||
""",
|
||||
)
|
||||
return True
|
||||
@@ -1,33 +0,0 @@
|
||||
class Insert:
|
||||
def __init__(self, connector, table_name: str, data: dict):
|
||||
self._connector = connector
|
||||
self._table_name = table_name
|
||||
self._columns = data.items()
|
||||
|
||||
def __gen_id(self) -> int:
|
||||
self._connector.execute(
|
||||
f"""
|
||||
SELECT MAX(id) FROM {self._table_name}
|
||||
""",
|
||||
)
|
||||
max_id = self._connector.fetchone(keys=["id"]).get("id")
|
||||
if max_id is None:
|
||||
return 1
|
||||
return max_id + 1
|
||||
|
||||
def one(self, data: dict):
|
||||
_id = self.__gen_id()
|
||||
data["id"] = _id
|
||||
columns = ", ".join(data.keys())
|
||||
placeholders = ", ".join([f":{key}" for key in data.keys()])
|
||||
self._connector.execute(
|
||||
f"""
|
||||
INSERT INTO {self._table_name} ({columns})
|
||||
VALUES ({placeholders})
|
||||
""",
|
||||
data,
|
||||
)
|
||||
|
||||
def many(self, data: list[dict]):
|
||||
for item in data:
|
||||
self.one(item)
|
||||
@@ -1,96 +0,0 @@
|
||||
class Select:
|
||||
def __init__(self, connector, table_name: str, data: dict):
|
||||
self._connector = connector
|
||||
self._table_name = table_name
|
||||
self._data = data
|
||||
|
||||
def id(self, _id: int):
|
||||
self._connector.execute(
|
||||
f"""
|
||||
SELECT * FROM {self._table_name}
|
||||
WHERE id = :id
|
||||
""",
|
||||
{"id": _id},
|
||||
)
|
||||
return self._connector.fetchone()
|
||||
|
||||
def all(self, limit: int = None):
|
||||
if limit is None:
|
||||
limit = 10000
|
||||
self._connector.execute(
|
||||
f"""
|
||||
SELECT * FROM {self._table_name} LIMIT {limit}
|
||||
""",
|
||||
)
|
||||
return self._connector.fetchall()
|
||||
|
||||
def one(
|
||||
self,
|
||||
keys: list[str] | None = None,
|
||||
conditions: dict = None,
|
||||
combine_operator: str = "AND",
|
||||
):
|
||||
if keys is None:
|
||||
columns = "*"
|
||||
else:
|
||||
columns = ", ".join(keys)
|
||||
condition_sql = self.__select_condition(conditions, combine_operator)
|
||||
self._connector.execute(
|
||||
f"""
|
||||
SELECT {columns} FROM {self._table_name}
|
||||
WHERE {condition_sql}
|
||||
""",
|
||||
conditions,
|
||||
)
|
||||
return self._connector.fetchone(keys)
|
||||
|
||||
def many(
|
||||
self,
|
||||
keys: list[str] | None = None,
|
||||
conditions: dict = None,
|
||||
combine_operator: str = "AND",
|
||||
limit: int = None,
|
||||
):
|
||||
if keys is None:
|
||||
columns = "*"
|
||||
else:
|
||||
columns = ", ".join(keys)
|
||||
if limit is None:
|
||||
limit = 10000
|
||||
condition_sql = self.__select_condition(conditions, combine_operator)
|
||||
self._connector.execute(
|
||||
f"""
|
||||
SELECT {columns} FROM {self._table_name}
|
||||
WHERE {condition_sql}
|
||||
LIMIT {limit}
|
||||
""",
|
||||
conditions,
|
||||
)
|
||||
return self._connector.fetchall(keys)
|
||||
|
||||
def column(self, keys: list[str]):
|
||||
columns = ", ".join(keys)
|
||||
self._connector.execute(
|
||||
f"""
|
||||
SELECT {columns} FROM {self._table_name}
|
||||
""",
|
||||
)
|
||||
return self._connector.fetchall(keys)
|
||||
|
||||
@staticmethod
|
||||
def __select_condition(conditions: dict, combine_operator: str = "AND"):
|
||||
if not conditions:
|
||||
raise ValueError("No conditions provided.")
|
||||
if combine_operator not in ["AND", "OR", "INSTR"]:
|
||||
raise ValueError(
|
||||
"Invalid combine_operator, must be 'AND' or 'OR' or 'INSTR'."
|
||||
)
|
||||
if combine_operator == "INSTR":
|
||||
condition_sql = f" AND ".join(
|
||||
[f"INSTR({key}, :{key})" for key in conditions.keys()]
|
||||
)
|
||||
else:
|
||||
condition_sql = f" {combine_operator} ".join(
|
||||
[f"{key} = :{key}" for key in conditions.keys()]
|
||||
)
|
||||
return condition_sql
|
||||
@@ -1,100 +0,0 @@
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Update:
|
||||
def __init__(self, connector, table_name: str, data: dict):
|
||||
self._connector = connector
|
||||
self._table_name = table_name
|
||||
self._example_data = data
|
||||
|
||||
def __table_exists(self) -> bool:
|
||||
self._connector.execute(
|
||||
f"""
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name='{self._table_name}'
|
||||
"""
|
||||
)
|
||||
return self._connector.fetch() is not None
|
||||
|
||||
def table(self):
|
||||
columns_list = [
|
||||
self.__python_to_sqlite_type(key, value)
|
||||
for key, value in self._example_data.items()
|
||||
]
|
||||
columns = ", ".join(columns_list)
|
||||
create_table_sql = f"CREATE TABLE IF NOT EXISTS {self._table_name} ({columns});"
|
||||
self._connector.execute(create_table_sql)
|
||||
logger.debug(f"Create table {self._table_name}.")
|
||||
self._connector.execute(f"PRAGMA table_info({self._table_name})")
|
||||
existing_columns = [x[1] for x in self._connector.fetch()]
|
||||
for key, value in self._example_data.items():
|
||||
if key not in existing_columns:
|
||||
insert_column = self.__python_to_sqlite_type(key, value)
|
||||
if value is None:
|
||||
value = "NULL"
|
||||
add_column_sql = f"ALTER TABLE {self._table_name} ADD COLUMN {insert_column} DEFAULT {value};"
|
||||
self._connector.execute(add_column_sql)
|
||||
logger.debug(f"Update table {self._table_name}.")
|
||||
|
||||
def one(self, data: dict) -> bool:
|
||||
_id = data["id"]
|
||||
set_sql = ", ".join([f"{key} = :{key}" for key in data.keys()])
|
||||
self._connector.execute(
|
||||
f"""
|
||||
UPDATE {self._table_name}
|
||||
SET {set_sql}
|
||||
WHERE id = :id
|
||||
""",
|
||||
data,
|
||||
)
|
||||
logger.debug(f"Update {_id} in {self._table_name}.")
|
||||
return True
|
||||
|
||||
def many(self, data: list[dict]) -> bool:
|
||||
columns = ", ".join([f"{key} = :{key}" for key in data[0].keys()])
|
||||
self._connector.executemany(
|
||||
f"""
|
||||
UPDATE {self._table_name}
|
||||
SET {columns}
|
||||
WHERE id = :id
|
||||
""",
|
||||
data,
|
||||
)
|
||||
logger.debug(f"Update {self._table_name}.")
|
||||
return True
|
||||
|
||||
def value(self, location: dict, set_value: dict) -> bool:
|
||||
set_sql = ", ".join([f"{key} = :{key}" for key in set_value.keys()])
|
||||
params = {**location, **set_value}
|
||||
self._connector.execute(
|
||||
f"""
|
||||
UPDATE {self._table_name}
|
||||
SET {set_sql}
|
||||
WHERE {location["key"]} = :{location["key"]}
|
||||
""",
|
||||
params,
|
||||
)
|
||||
logger.debug(f"Update {self._table_name}.")
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def __python_to_sqlite_type(key, value) -> str:
|
||||
if key == "id":
|
||||
column = "INTEGER PRIMARY KEY"
|
||||
elif isinstance(value, int):
|
||||
column = "INTEGER NOT NULL"
|
||||
elif isinstance(value, float):
|
||||
column = "REAL NOT NULL"
|
||||
elif isinstance(value, str):
|
||||
column = "TEXT NOT NULL"
|
||||
elif isinstance(value, bool):
|
||||
column = "INTEGER NOT NULL"
|
||||
elif isinstance(value, list):
|
||||
column = "TEXT NOT NULL"
|
||||
elif value is None:
|
||||
column = "TEXT"
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {type(value)}")
|
||||
return f"{key} {column}"
|
||||
@@ -2,72 +2,58 @@ import logging
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from module.database.connector import DataConnector
|
||||
from module.models.user import User
|
||||
from module.models.user import User, UserUpdate, UserLogin
|
||||
from module.security.jwt import get_password_hash, verify_password
|
||||
from module.database.engine import engine
|
||||
from sqlmodel import Session, select, SQLModel
|
||||
from sqlalchemy.exc import UnboundExecutionError, OperationalError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AuthDB(DataConnector):
|
||||
class UserDatabase(Session):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.__table_name = "user"
|
||||
if not self._table_exists(self.__table_name):
|
||||
self.__update_table()
|
||||
|
||||
def __update_table(self):
|
||||
db_data = self.__data_to_db(User())
|
||||
self._update_table(self.__table_name, db_data)
|
||||
self._insert(self.__table_name, db_data)
|
||||
|
||||
@staticmethod
|
||||
def __data_to_db(data: User) -> dict:
|
||||
db_data = data.dict()
|
||||
db_data["password"] = get_password_hash(db_data["password"])
|
||||
return db_data
|
||||
|
||||
@staticmethod
|
||||
def __db_to_data(db_data: dict) -> User:
|
||||
return User(**db_data)
|
||||
super().__init__(engine)
|
||||
statement = select(User)
|
||||
try:
|
||||
self.exec(statement)
|
||||
except OperationalError:
|
||||
SQLModel.metadata.create_all(engine)
|
||||
self.add(User())
|
||||
self.commit()
|
||||
|
||||
def get_user(self, username):
|
||||
self._cursor.execute(
|
||||
f"SELECT * FROM {self.__table_name} WHERE username=?", (username,)
|
||||
)
|
||||
result = self._cursor.fetchone()
|
||||
statement = select(User).where(User.username == username)
|
||||
result = self.exec(statement).first()
|
||||
if not result:
|
||||
return None
|
||||
db_data = dict(zip([x[0] for x in self._cursor.description], result))
|
||||
return self.__db_to_data(db_data)
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
return result
|
||||
|
||||
def auth_user(self, username, password) -> bool:
|
||||
self._cursor.execute(
|
||||
f"SELECT username, password FROM {self.__table_name} WHERE username=?",
|
||||
(username,),
|
||||
)
|
||||
result = self._cursor.fetchone()
|
||||
def auth_user(self, user: UserLogin) -> bool:
|
||||
statement = select(User).where(User.username == user.username)
|
||||
result = self.exec(statement).first()
|
||||
if not result:
|
||||
raise HTTPException(status_code=401, detail="User not found")
|
||||
if not verify_password(password, result[1]):
|
||||
if not verify_password(user.password, result.password):
|
||||
raise HTTPException(status_code=401, detail="Password error")
|
||||
return True
|
||||
|
||||
def update_user(self, username, update_user: User):
|
||||
def update_user(self, username, update_user: UserUpdate):
|
||||
# Update username and password
|
||||
new_username = update_user.username
|
||||
new_password = update_user.password
|
||||
self._cursor.execute(
|
||||
f"""
|
||||
UPDATE {self.__table_name}
|
||||
SET username = '{new_username}', password = '{get_password_hash(new_password)}'
|
||||
WHERE username = '{username}'
|
||||
"""
|
||||
)
|
||||
self._conn.commit()
|
||||
statement = select(User).where(User.username == username)
|
||||
result = self.exec(statement).first()
|
||||
if not result:
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
if update_user.username:
|
||||
result.username = update_user.username
|
||||
if update_user.password:
|
||||
result.password = get_password_hash(update_user.password)
|
||||
self.add(result)
|
||||
self.commit()
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
with AuthDB() as db:
|
||||
with UserDatabase() as db:
|
||||
# db.update_user(UserLogin(username="admin", password="adminadmin"), User(username="admin", password="cica1234"))
|
||||
db.update_user("admin", User(username="estrella", password="cica1234"))
|
||||
db.update_user("admin", UserUpdate(username="estrella", password="cica1234"))
|
||||
|
||||
Reference in New Issue
Block a user