Coverage for art_studio_tz/db_sql.py: 76%

1"""db_sql.py - SQLAlchemy based DB module""" 

2from sqlalchemy import create_engine, Column, Integer, String, DateTime, func, desc 

3from sqlalchemy.orm import sessionmaker, declarative_base 

4from sqlalchemy.exc import SQLAlchemyError 

5from typing import Any, Dict, List, Optional, Callable, cast 

6 

7Base = declarative_base() 

8 

9 

10class QuoteModel(Base): 

11 """SQLAlchemy model for the quotes table.""" 

12 __tablename__ = "quotes" 

13 

14 id = Column(Integer, primary_key=True, autoincrement=True) 

15 text = Column(String(1024), nullable=False) 

16 author = Column(String(255), nullable=True) 

17 timestep = Column(DateTime(timezone=True), server_default=func.now()) # время создания записи 

18 

19class DBsql: 

20 """SQLAlchemy based database handler 

21 """ 

22 def __init__(self, user: str, password: str, host: str, port: int, database: str): 

23 db_url = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}" 

24 self.engine = create_engine(db_url, echo=False, future=True) 

25 Base.metadata.create_all(self.engine) 

26 self.Session = sessionmaker(bind=self.engine) 

27 

28 def _execute(self, func: Callable, commit: bool = True, default=None): 

29 """ 

30 unified method to handle session lifecycle and errors 

31 """ 

32 session = self.Session() 

33 try: 

34 result = func(session) 

35 if commit: 

36 session.commit() 

37 return result 

38 except SQLAlchemyError as e: 

39 session.rollback() 

40 print(f"Database error: {e}") 

41 return default 

42 finally: 

43 session.close() 

44 

45 def create(self, item: Dict[str, Any]) -> Optional[int]: 

46 """create a new quote record 

47 

48 Args: 

49 item (Dict[str, Any]): quote data 

50 

51 Returns: 

52 Optional[int]: new record id or None if error 

53 """ 

54 def _create(session): 

55 quote = QuoteModel(**item) 

56 session.add(quote) 

57 return quote.id 

58 return self._execute(_create) 

59 

60 def read_all(self) -> List[Dict[str, Any]]: 

61 """Read all quote records 

62 

63 Returns: 

64 List[Dict[str, Any]]: list of dictionary with quote data 

65 """ 

66 def _read_all(session): 

67 quotes = session.query(QuoteModel).all() 

68 return [{c.name: getattr(q, c.name) for c in q.__table__.columns} for q in quotes] 

69 return cast(List[Dict[str, Any]], self._execute(_read_all, commit=False, default=[])) 

70 

71 def delete_all(self) -> bool: 

72 def _delete_all(session): 

73 session.query(QuoteModel).delete() 

74 return True 

75 return cast(bool, self._execute(_delete_all, default=False)) 

76 

77 def get_latest(self, n: int = 5) -> List[Dict[str, Any]]: 

78 """ Get latest n quotes 

79 

80 Args: 

81 n (int, optional): Number of quotes to get. Defaults to 5. 

82 

83 Returns: 

84 List[Dict[str, Any]]: list of dictionary with quote data 

85 """ 

86 def _latest(session): 

87 quotes = session.query(QuoteModel).order_by(desc(QuoteModel.timestep)).limit(n).all() 

88 return [{c.name: getattr(q, c.name) for c in q.__table__.columns} for q in quotes] 

89 return cast(List[Dict[str, Any]], self._execute(_latest, commit=False, default=[]))