Coverage for daklib/database/base.py: 97%

1# This program is free software; you can redistribute it and/or modify 

2# it under the terms of the GNU General Public License as published by 

3# the Free Software Foundation; either version 2 of the License, or 

4# (at your option) any later version. 

5 

6# This program is distributed in the hope that it will be useful, 

7# but WITHOUT ANY WARRANTY; without even the implied warranty of 

8# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

9# GNU General Public License for more details. 

10 

11# You should have received a copy of the GNU General Public License 

12# along with this program; if not, write to the Free Software 

13# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 

14 

15################################################################################ 

16 

17import datetime 

18 

19from sqlalchemy.event import listen 

20from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column 

21from sqlalchemy.schema import DDL, Table 

22from sqlalchemy.sql import func 

23from sqlalchemy.types import DateTime 

24 

25 

26class Base(DeclarativeBase): 

27 pass 

28 

29 

30class BaseMethods(Base): 

31 __abstract__ = True 

32 

33 @classmethod 

34 def get(cls, primary_key, session): 

35 """ 

36 This is a support function that allows getting an object by its primary 

37 key. 

38 

39 Architecture.get(3[, session]) 

40 

41 instead of the more verbose 

42 

43 session.query(Architecture).get(3) 

44 """ 

45 return session.query(cls).get(primary_key) 

46 

47 

48class BaseTimestamp(BaseMethods): 

49 __abstract__ = True 

50 

51 created: Mapped[datetime.datetime] = mapped_column( 

52 DateTime(timezone=True), nullable=False, server_default=func.now() 

53 ) 

54 modified: Mapped[datetime.datetime] = mapped_column( 

55 DateTime(timezone=True), nullable=False, server_default=func.now() 

56 ) 

57 

58 modified_trigger_function = DDL( 

59 """ 

60CREATE OR REPLACE FUNCTION tfunc_set_modified() RETURNS trigger 

61LANGUAGE plpgsql 

62AS $$ 

63 BEGIN NEW.modified = now(); return NEW; END; 

64$$ 

65 """ 

66 ) 

67 

68 modified_trigger = DDL( 

69 """ 

70CREATE TRIGGER %(table)s_modified BEFORE UPDATE ON %(fullname)s 

71FOR EACH ROW EXECUTE PROCEDURE tfunc_set_modified() 

72 """ 

73 ) 

74 

75 @classmethod 

76 def __table_cls__(cls, *arg, **kw): 

77 table = Table(*arg, **kw) 

78 listen( 

79 table, 

80 "after_create", 

81 cls.modified_trigger_function.execute_if(dialect="postgresql"), 

82 ) 

83 listen( 

84 table, 

85 "after_create", 

86 cls.modified_trigger.execute_if(dialect="postgresql"), 

87 ) 

88 return table