Coverage for anonymise/dbholder.py: 78%

1""" 

2crate_anon/anonymise/dbholder.py 

3 

4=============================================================================== 

5 

6 Copyright (C) 2015, University of Cambridge, Department of Psychiatry. 

7 Created by Rudolf Cardinal (rnc1001@cam.ac.uk). 

8 

9 This file is part of CRATE. 

10 

11 CRATE is free software: you can redistribute it and/or modify 

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

13 the Free Software Foundation, either version 3 of the License, or 

14 (at your option) any later version. 

15 

16 CRATE is distributed in the hope that it will be useful, 

17 but WITHOUT ANY WARRANTY; without even the implied warranty of 

18 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

19 GNU General Public License for more details. 

20 

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

22 along with CRATE. If not, see <https://www.gnu.org/licenses/>. 

23 

24=============================================================================== 

25 

26**Database "holder".** 

27 

28""" 

29 

30import logging 

31from typing import List, Optional, TYPE_CHECKING 

32 

33from sqlalchemy.engine import create_engine 

34from sqlalchemy.engine.base import Connection 

35from sqlalchemy.orm.session import sessionmaker, Session 

36from sqlalchemy.sql.schema import MetaData 

37 

38if TYPE_CHECKING: 

39 from crate_anon.anonymise.config import DatabaseSafeConfig 

40 

41log = logging.getLogger(__name__) 

42 

43 

44# ============================================================================= 

45# Convenience object 

46# ============================================================================= 

47 

48 

49class DatabaseHolder: 

50 """ 

51 Object to represent a connection to a database. 

52 """ 

53 

54 def __init__( 

55 self, 

56 name: str, 

57 url: str, 

58 srccfg: "DatabaseSafeConfig" = None, 

59 with_session: bool = False, 

60 with_conn: bool = True, 

61 reflect: bool = True, 

62 echo: bool = False, 

63 ) -> None: 

64 """ 

65 Args: 

66 name: internal database name 

67 url: SQLAlchemy URL 

68 srccfg: :class:`crate_anon.anonymise.config.DatabaseSafeConfig` 

69 with_session: create an SQLAlchemy Session? 

70 with_conn: create an SQLAlchemy connection (via an Engine)? 

71 reflect: read the database structure (when required)? 

72 echo: passed to SQLAlchemy's :func:`create_engine` 

73 """ 

74 self.name = name 

75 self.srccfg = srccfg 

76 self.engine = create_engine(url, echo=echo, future=True) 

77 self.conn = None # type: Optional[Connection] 

78 self.session = None # type: Optional[Session] 

79 self._reflect_on_request = reflect 

80 self._reflected = False 

81 self._table_names = [] # type: List[str] 

82 self._metadata = MetaData() 

83 log.debug(self.engine) # obscures password 

84 

85 if with_conn: # for raw connections 

86 self.conn = self.engine.connect() 

87 if with_session: # for ORM 

88 self.create_session() 

89 

90 def enable_reflect(self) -> None: 

91 """ 

92 Enables reflection, if it wasn't enabled to begin with. 

93 """ 

94 self._reflect_on_request = True 

95 

96 def create_session(self) -> None: 

97 """ 

98 Creates a database session, if not created to begin with. 

99 """ 

100 if not self.session: 

101 self.session = sessionmaker( 

102 bind=self.engine, future=True 

103 )() # type: Session 

104 

105 def _reflect(self) -> None: 

106 """ 

107 Perform the database reflection. 

108 

109 Reflection is expensive, so we defer unless required 

110 """ 

111 if not self._reflect_on_request: 

112 return 

113 log.info(f"Reflecting database: {self.name}") 

114 # self.table_names = get_table_names(self.engine) 

115 self._metadata.reflect(bind=self.engine, views=True) # include views 

116 self._table_names = [t.name for t in self._metadata.sorted_tables] 

117 self._reflected = True 

118 

119 def update_metadata(self) -> None: 

120 """ 

121 Updates the metadata, for example if a table has been dropped. 

122 """ 

123 self._metadata = MetaData() 

124 

125 @property 

126 def metadata(self) -> MetaData: 

127 """ 

128 Returns the SQLAlchemy :class:`MetaData`. If reflection is enabled, 

129 ensure the database has been reflected first. 

130 """ 

131 if not self._reflected: 

132 self._reflect() 

133 return self._metadata 

134 

135 @property 

136 def table_names(self) -> List[str]: 

137 """ 

138 Returns the table names from the database, if reflection is enabled. 

139 (Otherwise returns an empty list.) 

140 """ 

141 if not self._reflected: 

142 self._reflect() 

143 return self._table_names