Coverage for postrfp/shared/init/dbconfig.py: 98%

1import time 

2import logging 

3 

4from sqlalchemy.exc import OperationalError, ProgrammingError 

5from sqlalchemy import event, text 

6from sqlalchemy.engine import Connection, Engine, create_engine 

7 

8from postrfp import conf 

9from postrfp.shared.exceptions import TimezoneError 

10 

11log = logging.getLogger(__name__) 

12 

13 

14TEST_CONNECTION_STMT = "SELECT COUNT(*) FROM projects" 

15DELAY_INTERVAL = 10 

16 

17 

18def set_connection_timezone(dbapi_connection, _connrecord, tz: str = "+0:00"): 

19 """ 

20 Sets the 'sesson time_zone' for the given connection 

21 """ 

22 log.info("Connection setting timezone to UTC for db session") 

23 dbapi_connection.cursor().execute(f"SET session time_zone='{tz}' ") 

24 

25 

26def _check_connection_timezone(conn: Connection): 

27 """ 

28 Check that the session (connection) timezone is UTC 

29 """ 

30 tz_diff = conn.execute( 

31 text("SELECT TIME_TO_SEC(TIMEDIFF(NOW(),utc_timestamp()))") 

32 ).scalar_one() 

33 

34 if int(tz_diff) != 0: 

35 # This should never be able to happen, because the connection hook 

36 # sets the timezone to UTC. But if it does, raise an error. 

37 log.error(f"Timezone difference detected: {tz_diff} seconds. ") 

38 

39 session_tz = conn.execute(text("SELECT @@SESSION.TIME_ZONE")).scalar_one() 

40 global_tz_rows = conn.execute(text("SELECT @@GLOBAL.TIME_ZONE")).all() 

41 

42 global_tz = global_tz_rows[0][0] 

43 

44 raise TimezoneError( 

45 f"Timezone difference detected: {tz_diff} seconds. " 

46 f" Session Timezone: {session_tz}, Global (server) Timezone: {global_tz} " 

47 "Unable to proceed with the connection." 

48 ) 

49 

50 

51def check_connection(engine: Engine): 

52 """ 

53 Checks that: 

54 1 the database connection is live 

55 2 the database has a table named 'projects' 

56 3 the timezone of the DB session is the same as thaf of the local server 

57 """ 

58 for attempt in ("first", "second", "third"): 

59 log.info(f"{attempt} attempt at connecting to the DB") 

60 conn = None 

61 try: 

62 conn = engine.connect() 

63 conn.execute(text(TEST_CONNECTION_STMT)) 

64 

65 _check_connection_timezone(conn) 

66 

67 # Success: exit the loop. Finally will close the connection. 

68 break 

69 

70 except (ProgrammingError, OperationalError) as oe: 

71 if attempt == "third": 

72 if "unknown database" in str(oe).lower(): 

73 log.warning( 

74 f"Database doesn't exist at {engine.url}, quitting: {oe}" 

75 ) 

76 else: 

77 log.error( 

78 "Giving up database connection after three attempts: %s", oe 

79 ) 

80 raise oe 

81 

82 else: 

83 log.warning( 

84 f"DB Connection failed. Will try again after {DELAY_INTERVAL} seconds..." 

85 ) 

86 time.sleep(DELAY_INTERVAL) 

87 

88 finally: 

89 if conn is not None: 

90 conn.close() 

91 

92 

93def build_engine(echo=False) -> Engine: 

94 if conf.CONF is None: 

95 raise Exception( 

96 "--!!--postrfp not configured. Call sysconfig.configure_postrfp()--" 

97 ) 

98 

99 dsn = conf.CONF.sqlalchemy_dsn 

100 engine = create_engine( 

101 dsn, 

102 echo=echo, 

103 pool_recycle=3600, 

104 execution_options={"isolation_level": "REPEATABLE READ"}, 

105 ) 

106 

107 event.listen(engine, "connect", set_connection_timezone) 

108 

109 log.info("Created sqlalchemy engine with %s" % dsn) 

110 check_connection(engine) 

111 

112 return engine