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
« 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.
8:Platform: Linux/Windows | Python 3.10+
9:Developer: J Berendt
10:Email: support@s3dev.uk
12:Comments: This module contains *only* methods which can safely be
13 inherited and used by *any* of its subclasses.
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.
20 Any database-specific functionality must be contained in
21 that module.
23:Example:
25 For class-specific usage examples, please refer to the docstring
26 for the following classes:
28 - :class:`_DBIBase`
30"""
31# pylint: disable=wrong-import-order
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
41class SecurityWarning(Warning):
42 """Security warning stub-class."""
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.
50 Note:
51 This class is *not* designed to be interacted with directly.
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.
58 Args:
59 connstr (str): The database-specific SQLAlchemy connection
60 string.
62 :Example Use:
64 This low-level generalised class is designed to be inherited by
65 the calling/wrapping class as::
67 >>> from dblib.database import DBInterface
69 class MyDB(DBInterface):
71 def __init__(self, connstr: str):
72 super().__init__(connstr=('mysql+mysqlconnector://'
73 '<user>:<pwd>@<host>:<port>/'
74 '<db_name>'))
76 """
78 _PREFIX = '\n[DatabaseError]:'
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()
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
94 @property
95 def engine(self):
96 """Accessor to the ``sqlalchemy.engine.base.Engine`` object."""
97 return self._engine
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.
105 Important:
106 The following are *not* allowed to be executed by this
107 method:
109 - Statements containing multiple semi-colons (``;``).
110 - Statements containing a comment delimiter (``--``).
112 If found, a :class:`SecurityWarning` will be raised by the
113 :meth:`_is_dangerous` method.
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.
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.
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.
134 Warning:
136 1) Generally, whatever statement is passed into this method
137 **will be executed**, and may have *destructive
138 implications.*
140 2) This method contains a ``commit`` call.
142 If a statement is passed into this method, and the user has
143 the appropriate permissions - the change
144 **will be committed**.
146 **... HC SVNT DRACONES.**
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.
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.
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
181 def _create_engine(self) -> sa.engine.base.Engine:
182 """Create a database engine using the provided environment.
184 Returns:
185 sqlalchemy.engine.base.Engine: A sqlalchemy database engine
186 object.
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)
201 @staticmethod
202 def _is_dangerous(stmt: str) -> bool:
203 """Perform a dirty security check for injection attempts.
205 Args:
206 stmt (str): SQL statement to be potentially executed.
208 Raises:
209 SecurityWarning: If there are multiple semi-colons (``;``)
210 in the statement, or any comment delimiters (``--``).
212 Returns:
213 bool: False if the checks pass.
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
226 def _report_sa_error(self, msg: str, error: SQLAlchemyError): # pragma: nocover
227 """Report SQLAlchemy error to the terminal.
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.
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)
243 @staticmethod
244 def _result_to_df__cursor(result: sa.engine.cursor.CursorResult) -> pd.DataFrame:
245 """Convert a ``CursorResult`` object to a DataFrame.
247 If the cursor did not return results, an empty DataFrame
248 containing the column names only, is returned.
250 Args:
251 result (sqlalchemy.engine.cursor.CursorResult): Object to
252 be converted.
254 Returns:
255 pd.DataFrame: A ``pandas.DataFrame`` object containing the
256 cursor's data.
258 """
259 return pd.DataFrame(result, columns=result.keys())
261 @staticmethod
262 def _result_to_df__stored(result: object) -> pd.DataFrame:
263 """Convert a ``MySQLCursor.stored_results`` object to a DataFrame.
265 Args:
266 result (object): The ``cursor.stored_results()`` object from
267 a ``sqlalchemy`` or ``mysql.connector`` procedure call.
269 Returns:
270 pd.DataFrame: A DataFrame containing the results from the
271 procedure call.
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