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

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

33 import pandas as pd

34 import traceback

35 import sqlalchemy as sa

36 from sqlalchemy.exc import SQLAlchemyError

37 from utils4.reporterror import reporterror

38 from utils4.user_interface import ui

39

40

41 class SecurityWarning(Warning):

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

43

44

45 class _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