Coverage for /var/devmt/py/dbilib_0.0.1.dev1/dbilib/_dbi_base.py: 100%

54 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-08-07 23:49 +0100

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4:Purpose: This module contains the library's *base* database methods 

5 and attribute accessors, which are designed to be 

6 specialised by the database-specific modules and classes. 

7 

8:Platform: Linux/Windows | Python 3.10+ 

9:Developer: J Berendt 

10:Email: support@s3dev.uk 

11 

12:Comments: This module contains *only* methods which can safely be 

13 inherited and used by *any* of its subclasses. 

14 

15 In other words, this module should *not* contain any import 

16 statement, or uses of these imports, which if used in a 

17 database-specific module will cause a crash due to a missing 

18 library. 

19 

20 Any database-specific functionality must be contained in 

21 that module. 

22 

23:Example: 

24 

25 For class-specific usage examples, please refer to the docstring 

26 for the following classes: 

27 

28 - :class:`_DBIBase` 

29 

30""" 

31# pylint: disable=wrong-import-order 

32 

33import pandas as pd 

34import traceback 

35import sqlalchemy as sa 

36from sqlalchemy.exc import SQLAlchemyError 

37from utils4.reporterror import reporterror 

38from utils4.user_interface import ui 

39 

40 

41class SecurityWarning(Warning): 

42 """Security warning stub-class.""" 

43 

44 

45class _DBIBase: 

46 """This class holds the methods and properties which are used across 

47 all databases. Each of the database-specific constructors inherits 

48 this class for its members. 

49 

50 Note: 

51 This class is *not* designed to be interacted with directly. 

52 

53 Rather, please use the :class:`database.DBInterface` class 

54 instead, as the proper interface class has an automatic switch 

55 for database interfaces, based on the ``sqlalchemy.Engine`` 

56 object which is created from the connection string. 

57 

58 Args: 

59 connstr (str): The database-specific SQLAlchemy connection 

60 string. 

61 

62 :Example Use: 

63 

64 This low-level generalised class is designed to be inherited by 

65 the calling/wrapping class as:: 

66 

67 >>> from dblib.database import DBInterface 

68 

69 class MyDB(DBInterface): 

70 

71 def __init__(self, connstr: str): 

72 super().__init__(connstr=('mysql+mysqlconnector://' 

73 '<user>:<pwd>@<host>:<port>/' 

74 '<db_name>')) 

75 

76 """ 

77 

78 _PREFIX = '\n[DatabaseError]:' 

79 

80 def __init__(self, connstr: str): 

81 """Class initialiser.""" 

82 self._connstr = connstr 

83 self._engine = None 

84 if connstr: 

85 # Testing: Enable an instance to be created without a 

86 # connection string. 

87 self._engine = self._create_engine() 

88 

89 @property 

90 def database_name(self): 

91 """Accessor to the database name used by the :attr:`engine` object.""" 

92 return self._engine.url.database 

93 

94 @property 

95 def engine(self): 

96 """Accessor to the ``sqlalchemy.engine.base.Engine`` object.""" 

97 return self._engine 

98 

99 def execute_query(self, 

100 stmt: str, 

101 params: dict=None, 

102 raw: bool=True) -> list | pd.DataFrame | None: 

103 """Execute a query statement. 

104 

105 Important: 

106 The following are *not* allowed to be executed by this 

107 method: 

108 

109 - Statements containing multiple semi-colons (``;``). 

110 - Statements containing a comment delimiter (``--``). 

111 

112 If found, a :class:`SecurityWarning` will be raised by the 

113 :meth:`_is_dangerous` method. 

114 

115 Args: 

116 stmt (str): Statement to be executed. The parameter bindings 

117 are to be written in colon format. 

118 params (dict, optional): Parameter key/value bindings as a 

119 dictionary, if applicable. Defaults to None. 

120 raw (bool, optional): Return the data in 'raw' (tuple) 

121 format rather than as a formatted DataFrame. 

122 Defaults to True for efficiency. 

123 

124 If the query did not return results and the ``raw`` argument is 

125 False, an empty DataFrame containing the column names only, is 

126 returned. 

127 

128 Note: 

129 In the SQL query, the bind parameters are specified by name, 

130 using the format ``:bind_name``. The ``params`` dictionary 

131 argument must contain the associated parameter name/value 

132 bindings. 

133 

134 Warning: 

135 

136 1) Generally, whatever statement is passed into this method 

137 **will be executed**, and may have *destructive 

138 implications.* 

139 

140 2) This method contains a ``commit`` call. 

141 

142 If a statement is passed into this method, and the user has 

143 the appropriate permissions - the change 

144 **will be committed**. 

145 

146 **... HC SVNT DRACONES.** 

147 

148 Returns: 

149 list | pd.DataFrame | None: If the ``raw`` parameter is 

150 True, a list of tuples containing values is returned. 

151 Otherwise, a ``pandas.DataFrame`` object containing the 

152 returned data is returned. 

153 

154 If this method is called with a script which does not return 

155 results, for example a CREATE script, None is returned; 

156 regardless of the value passed to the ``raw`` parameter. 

157 

158 """ 

159 # Additional else and return used for clarity. 

160 # pylint: disable=no-else-return 

161 # The error does have a _message member. 

162 # pylint: disable=no-member 

163 try: 

164 # Perform a cursory 'security check.' 

165 if not self._is_dangerous(stmt=stmt): 

166 with self._engine.connect() as conn: 

167 result = conn.execute(sa.text(stmt), params) 

168 conn.commit() 

169 conn.close() 

170 if raw: 

171 return result.fetchall() 

172 else: 

173 return self._result_to_df__cursor(result=result) 

174 except SecurityWarning: 

175 print(traceback.format_exc()) 

176 except Exception as err: 

177 if 'object does not return rows' not in err._message(): 

178 reporterror(err) 

179 return None 

180 

181 def _create_engine(self) -> sa.engine.base.Engine: 

182 """Create a database engine using the provided environment. 

183 

184 Returns: 

185 sqlalchemy.engine.base.Engine: A sqlalchemy database engine 

186 object. 

187 

188 """ 

189 # ???: Do these values need to be moved to an external config? 

190 # Added in s3ddb v0.7.0: 

191 # The pool_* arguments to prevent MySQL timeout which causes 

192 # a broken pipe and lost connection errors. 

193 return sa.create_engine(url=self._connstr, 

194 poolclass=sa.pool.QueuePool, 

195 pool_size=20, 

196 pool_recycle=3600, 

197 pool_timeout=30, 

198 pool_pre_ping=True, 

199 max_overflow=0) 

200 

201 @staticmethod 

202 def _is_dangerous(stmt: str) -> bool: 

203 """Perform a dirty security check for injection attempts. 

204 

205 Args: 

206 stmt (str): SQL statement to be potentially executed. 

207 

208 Raises: 

209 SecurityWarning: If there are multiple semi-colons (``;``) 

210 in the statement, or any comment delimiters (``--``). 

211 

212 Returns: 

213 bool: False if the checks pass. 

214 

215 """ 

216 # import sys 

217 if stmt.count(';') > 1: 

218 msg = 'Multiple statements are disallowed for security reasons.' 

219 raise SecurityWarning(msg) 

220 # sys.exit(1) 

221 if '--' in stmt: 

222 msg = 'Comments are not allowed in the statement for security reasons.' 

223 raise SecurityWarning(msg) 

224 return False 

225 

226 def _report_sa_error(self, msg: str, error: SQLAlchemyError): # pragma: nocover 

227 """Report SQLAlchemy error to the terminal. 

228 

229 Args: 

230 msg (str): Additional error to be displayed. This message 

231 will be automatically prefixed with '[DatabaseError]: ' 

232 error (sqlalchemy.exc.SQLAlchemyError): Caught error object 

233 from the try/except block. 

234 

235 """ 

236 msg = f'\n{self._PREFIX} {msg}' 

237 stmt = f'- Statement: {error.statement}' 

238 errr = f'- Error: {str(error.orig)}' 

239 ui.print_alert(text=msg) 

240 ui.print_alert(text=stmt) 

241 ui.print_alert(text=errr) 

242 

243 @staticmethod 

244 def _result_to_df__cursor(result: sa.engine.cursor.CursorResult) -> pd.DataFrame: 

245 """Convert a ``CursorResult`` object to a DataFrame. 

246 

247 If the cursor did not return results, an empty DataFrame 

248 containing the column names only, is returned. 

249 

250 Args: 

251 result (sqlalchemy.engine.cursor.CursorResult): Object to 

252 be converted. 

253 

254 Returns: 

255 pd.DataFrame: A ``pandas.DataFrame`` object containing the 

256 cursor's data. 

257 

258 """ 

259 return pd.DataFrame(result, columns=result.keys()) 

260 

261 @staticmethod 

262 def _result_to_df__stored(result: object) -> pd.DataFrame: 

263 """Convert a ``MySQLCursor.stored_results`` object to a DataFrame. 

264 

265 Args: 

266 result (object): The ``cursor.stored_results()`` object from 

267 a ``sqlalchemy`` or ``mysql.connector`` procedure call. 

268 

269 Returns: 

270 pd.DataFrame: A DataFrame containing the results from the 

271 procedure call. 

272 

273 """ 

274 df = pd.DataFrame() 

275 try: 

276 # There is only one item in the iterable. 

277 # However, if the iterable is empty, a StopIteration error is raised 

278 # when using x = next(result); so a loop is used instead. 

279 for x in result: 

280 df = pd.DataFrame(data=x.fetchall(), columns=x.column_names) 

281 except Exception as err: 

282 reporterror(err) 

283 return df