Coverage for C:\var\devmt\py\dbilib_0.5.2\dbilib\_dbi_mssql.py: 100%

108 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2026-03-02 11:16 +0000

1#!/usr/bin/env python3 

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

3""" 

4:Purpose: This module contains the library's *Microsoft SQL Server* 

5 database methods and attribute accessors; which are a 

6 specialised version of the :class:`_dbi_base._DBIBase` class 

7 methods. 

8 

9:Platform: Linux/Windows | Python 3.10+ 

10:Developer: J Berendt, J Preston 

11:Email: support@s3dev.uk 

12 

13:Comments: n/a 

14 

15:Example: 

16 

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

18 for the following classes: 

19 

20 - :class:`_DBIMSSQL` 

21 

22""" 

23# pylint: disable=import-error 

24 

25from __future__ import annotations 

26 

27import sqlalchemy as sa 

28from sqlalchemy.exc import SQLAlchemyError 

29from utils4.reporterror import reporterror 

30from utils4.user_interface import ui 

31# locals 

32try: 

33 from ._dbi_base import _DBIBase, ExitCode 

34except ImportError: 

35 from _dbi_base import _DBIBase, ExitCode 

36 

37 

38class _DBIMSSQL(_DBIBase): 

39 """This *private* class holds the methods and properties which are 

40 used for accessing Microsoft SQL Server databases. 

41 

42 Note: 

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

44 

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

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

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

48 object which is created from the connection string. 

49 

50 Args: 

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

52 string. 

53 

54 :Example Use: 

55 

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

57 the calling/wrapping class as:: 

58 

59 >>> from dbilib.database import DBInterface 

60 

61 class MyDB(DBInterface): 

62 

63 def __init__(self, connstr: str): 

64 super().__init__(connstr=('mssql+pyodbc://' 

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

66 '<db_name>')) 

67 

68 

69 'Userless' trusted connections can be used in the connection 

70 string as follows:: 

71 

72 >>> from dbilib.database import DBInterface 

73 

74 class MyDB(DBInterface): 

75 

76 def __init__(self, connstr: str): 

77 super().__init__(connstr=('mssql+pyodbc' 

78 '://:@<host>:<port>/' 

79 '<dbname>' 

80 '?driver=<driver-name>' 

81 '&trusted_connection=yes')) 

82 

83 """ 

84 

85 # The __init__ method is implemented in the parent class. 

86 

87 def backup(self, table_name: str, verbose: bool=True) -> ExitCode: 

88 """Backup the given table to the backup database. 

89 

90 Args: 

91 table_name (str): Name of the table to be backed up. 

92 verbose (bool, optional): Display helpful text indicating the 

93 status of the backup. Defaults to True. 

94 

95 .. important:: 

96 The backup database (which is implicitly determined by the 

97 engine's database name) **must** exist, or this method will 

98 fail. See below for the backup database's naming convention. 

99 

100 Due to MSSQL's (interesting) handling of the 

101 ``CREATE DATABASE`` statement, the database cannot be created 

102 by this method for you; sorry. Cheers MS! 

103 

104 

105 .. note:: 

106 

107 The backup database name is derived by prepending 

108 ``'__bak__'`` to the database name. 

109 

110 This obfuscation was done intentionally to help prevent a 

111 user from click-selecting the wrong database (in SSMS) by 

112 accident. 

113 

114 Returns: 

115 ExitCode: The exit code enumerator object associated to the 

116 status of the backup process. 

117 

118 """ 

119 bkdb = f'__bak__{self.database_name}' 

120 s1, s2, s3 = False, False, False 

121 s1 = self.table_exists(table_name=table_name, verbose=verbose) 

122 if s1: s2 = self.database_exists(database_name=bkdb, verbose=verbose) 

123 if s2: s3 = self._backup(table_name=table_name, bkdb_name=bkdb) 

124 if verbose: self._print_summary(success=all((s1, s2, s3))) 

125 if not s1: return ExitCode.ERR_BKUP_TBNEX 

126 if not s2: return ExitCode.ERR_BKUP_DBNEX 

127 if not s3: return ExitCode.ERR_BKUP_CKSUM 

128 return ExitCode.OK 

129 

130 # pylint: disable=line-too-long 

131 def call_procedure(self, 

132 proc: str, 

133 *, 

134 params: dict | tuple=None, 

135 schema: str='dbo', 

136 paramnames: list | tuple=None, 

137 raw: bool=True, 

138 return_status: bool=False) -> pd.DataFrame | tuple[pd.DataFrame | tuple, bool]: # noqa # pylint: disable=undefined-variable 

139 """Call a stored procedure, and return as a DataFrame. 

140 

141 Args: 

142 proc (str): Name of the stored procedure to call. 

143 params (dict): A dictionary containing the parameter values 

144 to be used as key/value pairs, with the keys being the 

145 parameter names, and values being the data values passed 

146 into the procedure. 

147 paramnames (list|tuple, optional): An iterable object 

148 containing the procedure's parameter names, in order. 

149 Defaults to None. If not provided, these are collected 

150 from the database via the :meth:`get_parameter_names` 

151 method. For efficiency, the parameter names should be 

152 passed if making repeated calls to the procedure. 

153 schema (str, optional): Name of the schema to use. 

154 Defaults to dbo. 

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

156 rather than as a formatted DataFrame. Defaults to True 

157 for efficiency. 

158 return_status (bool, optional): Return the method's success 

159 status. Defaults to False. 

160 

161 Returns: 

162 pd.DataFrame | tuple[pd.DataFrame | tuple, bool]: 

163 If the ``return_status`` argument is True, a tuple of the 

164 data and the method's return status is returned as:: 

165 

166 (data, status) 

167 

168 Otherwise, only the data is returned. 

169 

170 """ 

171 # pylint: disable=consider-using-f-string # No, need the formatter. 

172 data = None 

173 success = False 

174 try: 

175 # Collect parameter names for the EXEC call if not provided. 

176 if not paramnames: 

177 paramnames = self.get_parameter_names(proc=proc, schema=schema) 

178 paramdef = ', '.join(map(':{}'.format, paramnames)) 

179 with self.engine.connect() as con: 

180 resp = con.execute(sa.text(f'EXEC {schema}.{proc} {paramdef}'), params) 

181 if resp.returns_rows: 

182 if raw: 

183 data = resp.fetchall() 

184 success = bool(data) 

185 else: 

186 data = self._result_to_df__cursor(result=resp) 

187 success = not data.empty 

188 con.close() 

189 except SQLAlchemyError as err: 

190 msg = f'Error occurred while running the USP: {proc}.' 

191 self._report_sa_error(msg=msg, error=err) 

192 except Exception as err: 

193 reporterror(error=err) 

194 return (data, success) if return_status else data 

195 

196 def call_procedure_update(self, 

197 proc: str, 

198 *, 

199 data: dict, 

200 paramnames: list | tuple=None, 

201 schema: str='dbo', 

202 return_id: bool=False) -> bool | tuple: 

203 """Call an *update* or *insert* stored procedure. 

204 

205 Note: 

206 Results are *not* returned from this call, only a boolean 

207 status flag and the optional last row ID, which **must be 

208 provided by the USP** if desired. 

209 

210 If results are desired, please use the 

211 :meth:`~call_procedure` method. 

212 

213 Args: 

214 proc (str): Name of the stored procedure to call. 

215 data (dict | list[dict]): A dictionary containing the data to be loaded 

216 as key/value pairs, with the keys being the parameter 

217 names, and values being the data values. 

218 

219 Note: 

220 If *multiple* rows are to be inserted, this 

221 parameter must contain a list of dictionaries, with 

222 each row to be added as a dictionary of it's own. 

223 

224 paramnames (list|tuple, optional): An iterable object 

225 containing the procedure's parameter names, in order. 

226 Defaults to None. If not provided, these are collected 

227 from the database via the :meth:`get_parameter_names` 

228 method. For efficiency, the parameter names should be 

229 passed if making repeated calls to the procedure. 

230 schema (str, optional): Name of the schema to use. 

231 Defaults to dbo. 

232 return_id (bool, optional): Return the ID of the last 

233 inserted row. If a duplicate constraint is encountered, 

234 -1 is returned as the row ID. **See note above.** 

235 Defaults to False. 

236 

237 Returns: 

238 bool | tuple: If ``return_id`` is False, True is 

239 returned if the procedure completed successfully, otherwise 

240 False. If ``return_id`` is True, a tuple containing the 

241 ID of the last inserted row (or -1 on duplicate) and the 

242 execution success flag are returned as:: 

243 

244 (rowid, success_flag) 

245 

246 """ 

247 # pylint: disable=consider-using-f-string # No, need the formatter. 

248 try: 

249 success = False 

250 rowid = None 

251 # Collect parameter names for the EXEC call if not provided. 

252 if not paramnames: 

253 paramnames = self.get_parameter_names(proc=proc, schema=schema) 

254 params = ', '.join(map('@{} = :{}'.format, paramnames, paramnames)) 

255 with self.engine.connect() as con: 

256 resp = con.execute(sa.text(f'EXEC {schema}.{proc} {params}'), data) 

257 if resp.returns_rows: 

258 rowid = resp.fetchall() 

259 con.commit() 

260 con.close() 

261 success = True 

262 except Exception as err: 

263 if 'Cannot insert duplicate key' in repr(err): 

264 msg = f'{self._PREFIXW.strip()} Duplicate record detected, skipping.' 

265 rowid = [(-1,)] # Match format of a returned row ID. 

266 ui.print_warning(text=msg) 

267 else: 

268 reporterror(err) 

269 return (rowid, success) if return_id else success 

270 

271 def call_procedure_update_raw(self, 

272 proc: str, 

273 *, 

274 data: dict, 

275 paramnames: list | tuple=None) -> None: 

276 """Call an *update* or *insert* stored procedure, without error 

277 handling. 

278 

279 .. warning:: 

280 This method is **unprotected**, perhaps use 

281 :meth:`~call_procedure_update` instead. 

282 

283 This 'raw' method *does not* contain an error handler. It is 

284 (by design) the responsibility of the caller to contain and 

285 control the errors. 

286 

287 The purpose of this raw method is to enable the caller method to 

288 contain and control the errors which might be generated from a 

289 USP call, for example a **duplicate key** error. 

290 

291 Args: 

292 proc (str): Name of the stored procedure to call. 

293 data (dict): A dictionary containing the data to be loaded 

294 as key/value pairs, with the keys being the parameter 

295 names, and values being the data values. 

296 paramnames (list|tuple, optional): An iterable object 

297 containing the procedure's parameter names, in order. 

298 Defaults to None. If not provided, these are collected 

299 from the database via the :meth:`get_parameter_names` 

300 method. For efficiency, the parameter names should be 

301 passed if making repeated calls to the procedure. 

302 

303 """ 

304 # pylint: disable=consider-using-f-string # No, need the formatter. 

305 # Collect parameter names for the EXEC call if not provided. 

306 if not paramnames: 

307 paramnames = self.get_parameter_names(proc=proc) 

308 params = ', '.join(map(':{}'.format, paramnames)) 

309 with self.engine.connect() as con: 

310 _ = con.execute(sa.text(f'EXEC {proc} {params}'), data) 

311 con.commit() 

312 con.close() 

313 

314 def checksum(self, 

315 table_name: str, 

316 database_name: str=None, 

317 schema: str='dbo') -> int | None: 

318 """Calculate a hash (checksum) on the given table. 

319 

320 Args: 

321 table_name (str) Name of the table against which the checksum 

322 is to be calculated. 

323 database_name (str) Name of the database to use. 

324 This argument can be used if the table resides in a 

325 different database than the one to which the engine 

326 object already points. Defaults to None. 

327 schema (str, optional): Name of the schema to use. 

328 Defaults to dbo. 

329 

330 This method wraps the ``CHECKSUM_AGG`` and ``BINARY_CHECKSUM`` 

331 MSSQL functions. 

332 

333 Returns: 

334 int | None: A signed integer representation of the table's 

335 hash value, if the table exists. Otherwise, None. 

336 

337 """ 

338 db = database_name if database_name else self.database_name 

339 stmt = f'SELECT CHECKSUM_AGG(BINARY_CHECKSUM(*)) FROM [{db}].[{schema}].[{table_name}]' 

340 if all((not self._is_dangerous(stmt=stmt), 

341 self.table_exists(table_name=table_name, database_name=db))): 

342 return self.execute_query(stmt)[0][0] 

343 return None 

344 

345 def database_exists(self, database_name: str, verbose: bool=False) -> bool: 

346 """Using the ``engine`` object, test if the given database exists. 

347 

348 Args: 

349 database_name (str): Name of the database to test. 

350 verbose (bool, optional): Print a message if the database 

351 does not exist. Defaults to False. 

352 

353 Returns: 

354 bool: True if the given database exists, otherwise False. 

355 

356 """ 

357 exists = False 

358 stmt = f'select count(*) from [sys].[databases] where [name] = \'{database_name}\'' 

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

360 exists = bool(self.execute_query(stmt, raw=True)[0][0]) 

361 if (not exists) & verbose: 

362 msg = f'Database does not exist: {database_name}' 

363 ui.print_warning(text=msg) 

364 return exists 

365 

366 def get_parameter_names(self, proc: str, schema: str='dbo') -> tuple: 

367 """Retrieve the parameter names for the given USP. 

368 

369 For portability, this method has been updated to use an embedded 

370 query rather than a USP. 

371 

372 Args: 

373 proc (str): Name of the target stored procedure. 

374 schema (str, optional): Name of the schema to use. 

375 Defaults to dbo. 

376 Returns: 

377 tuple: A tuple of parameter names for the given USP. 

378 

379 """ 

380 # TODO: Make this more robust by removing the hardcoded SUBSTRING and 

381 # test if the first char is '@' before removing. 

382 stmt = f""" 

383 /* Use SUBSTRING to remove the '@' from the parameter name. */ 

384 SELECT 

385 SUBSTRING([name], 2, 99) AS parameters 

386 FROM 

387 [sys].[parameters] 

388 WHERE 

389 [object_id] = ( SELECT [object_id] FROM [sys].[objects] WHERE [name] = '{proc}' 

390 AND schema_id = SCHEMA_ID('{schema}') ) 

391 ORDER BY 

392 [parameter_id]; 

393 """ 

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

395 resp = self.execute_query(stmt=stmt, raw=True) 

396 if resp: 

397 return next(zip(*resp)) 

398 return () 

399 return () # pragma: nocover # Unreachable 

400 

401 def table_exists(self, 

402 table_name: str, 

403 database_name: str=None, 

404 schema: str='dbo', 

405 verbose: bool=False) -> bool: 

406 """Using the ``engine`` object, test if the given table exists. 

407 

408 Args: 

409 table_name (str): Name of the table to test. 

410 database_name (str) Name of the database to use. 

411 This argument can be used if the table resides in a 

412 different database than the one to which the engine 

413 object already points. Defaults to None. 

414 schema (str, optional): Name of the schema to use. 

415 Defaults to dbo. 

416 verbose (bool, optional): Print a message if the table does 

417 not exist. Defaults to False. 

418 

419 Returns: 

420 bool: True if the given table exists, otherwise False. 

421 

422 """ 

423 exists = False 

424 db = database_name if database_name else self.database_name 

425 params = { 

426 'table_catalog': db, 

427 'table_name': table_name, 

428 'table_schema': schema, 

429 } 

430 stmt = (f'select count(*) from [{db}].[information_schema].[tables] ' 

431 'where table_catalog = :table_catalog ' 

432 'and table_schema = :table_schema ' 

433 'and table_name = :table_name') 

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

435 exists = bool(self.execute_query(stmt, params=params, raw=True)[0][0]) 

436 if (not exists) & verbose: 

437 msg = f'Table does not exist: {db}.{table_name}' 

438 ui.print_warning(text=msg) 

439 return exists 

440 

441 def _backup(self, table_name: str, bkdb_name: str, schema: str='dbo') -> bool: 

442 """Perform the table backup to the backup database. 

443 

444 Args: 

445 table_name (str): Name of the table to be backed up. 

446 bkdb_name (str): Name of the backup database. 

447 schema (str, optional): Name of the schema to use. 

448 Defaults to dbo. 

449 Returns: 

450 bool: True if the backup was successful, otherwise False. 

451 A successful backup is determined by verifying matching table 

452 checksum values between the origin and backup tables. 

453 

454 """ 

455 stmt1 = f'DROP TABLE IF EXISTS [{bkdb_name}].[{schema}].[{table_name}]' 

456 stmt2 = f'SELECT * INTO [{bkdb_name}].[{schema}].[{table_name}] FROM [{schema}].[{table_name}]' 

457 if all((not self._is_dangerous(stmt=stmt1), not self._is_dangerous(stmt=stmt2))): 

458 self.execute_query(stmt1) 

459 self.execute_query(stmt2) 

460 ck1 = self.checksum(table_name=table_name, database_name=self.database_name, schema=schema) 

461 ck2 = self.checksum(table_name=table_name, database_name=bkdb_name, schema=schema) 

462 return ck1 == ck2 

463 

464 @staticmethod 

465 def _print_summary(success: bool) -> None: 

466 """Print a short end-of-processing summary. 

467 

468 Args: 

469 success (bool): Flag indicating if the backup was successful. 

470 

471 This message is designed to be short and concise, as the backup 

472 is designed to be called by other applications, so a short 

473 message is preferable. 

474 

475 """ 

476 if success: 

477 ui.print_normal('Table backup successful.') 

478 else: 

479 ui.print_warning('Table backup failed.')