Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/backends/base/base.py: 67%

1 import _thread

2 import copy

3 import datetime

4 import logging

5 import threading

6 import time

7 import warnings

8 import zoneinfo

9 from collections import deque

10 from contextlib import contextmanager

11

12 from plain.models.backends import utils

13 from plain.models.backends.base.validation import BaseDatabaseValidation

14 from plain.models.backends.signals import connection_created

15 from plain.models.backends.utils import debug_transaction

16 from plain.models.db import (

17 DEFAULT_DB_ALIAS,

18 DatabaseError,

19 DatabaseErrorWrapper,

20 NotSupportedError,

21)

22 from plain.models.transaction import TransactionManagementError

23 from plain.runtime import settings

24 from plain.utils.functional import cached_property

25

26 NO_DB_ALIAS = "__no_db__"

27 RAN_DB_VERSION_CHECK = set()

28

29 logger = logging.getLogger("plain.models.backends.base")

30

31

32 class BaseDatabaseWrapper:

33 """Represent a database connection."""

34

35 # Mapping of Field objects to their column types.

36 data_types = {}

37 # Mapping of Field objects to their SQL suffix such as AUTOINCREMENT.

38 data_types_suffix = {}

39 # Mapping of Field objects to their SQL for CHECK constraints.

40 data_type_check_constraints = {}

41 ops = None

42 vendor = "unknown"

43 display_name = "unknown"

44 SchemaEditorClass = None

45 # Classes instantiated in __init__().

46 client_class = None

47 creation_class = None

48 features_class = None

49 introspection_class = None

50 ops_class = None

51 validation_class = BaseDatabaseValidation

52

53 queries_limit = 9000

54

55 def __init__(self, settings_dict, alias=DEFAULT_DB_ALIAS):

56 # Connection related attributes.

57 # The underlying database connection.

58 self.connection = None

59 # `settings_dict` should be a dictionary containing keys such as

60 # NAME, USER, etc. It's called `settings_dict` instead of `settings`

61 # to disambiguate it from Plain settings modules.

62 self.settings_dict = settings_dict

63 self.alias = alias

64 # Query logging in debug mode or when explicitly enabled.

65 self.queries_log = deque(maxlen=self.queries_limit)

66 self.force_debug_cursor = False

67

68 # Transaction related attributes.

69 # Tracks if the connection is in autocommit mode. Per PEP 249, by

70 # default, it isn't.

71 self.autocommit = False

72 # Tracks if the connection is in a transaction managed by 'atomic'.

73 self.in_atomic_block = False

74 # Increment to generate unique savepoint ids.

75 self.savepoint_state = 0

76 # List of savepoints created by 'atomic'.

77 self.savepoint_ids = []

78 # Stack of active 'atomic' blocks.

79 self.atomic_blocks = []

80 # Tracks if the outermost 'atomic' block should commit on exit,

81 # ie. if autocommit was active on entry.

82 self.commit_on_exit = True

83 # Tracks if the transaction should be rolled back to the next

84 # available savepoint because of an exception in an inner block.

85 self.needs_rollback = False

86 self.rollback_exc = None

87

88 # Connection termination related attributes.

89 self.close_at = None

90 self.closed_in_transaction = False

91 self.errors_occurred = False

92 self.health_check_enabled = False

93 self.health_check_done = False

94

95 # Thread-safety related attributes.

96 self._thread_sharing_lock = threading.Lock()

97 self._thread_sharing_count = 0

98 self._thread_ident = _thread.get_ident()

99

100 # A list of no-argument functions to run when the transaction commits.

101 # Each entry is an (sids, func, robust) tuple, where sids is a set of

102 # the active savepoint IDs when this function was registered and robust

103 # specifies whether it's allowed for the function to fail.

104 self.run_on_commit = []

105

106 # Should we run the on-commit hooks the next time set_autocommit(True)

107 # is called?

108 self.run_commit_hooks_on_set_autocommit_on = False

109

110 # A stack of wrappers to be invoked around execute()/executemany()

111 # calls. Each entry is a function taking five arguments: execute, sql,

112 # params, many, and context. It's the function's responsibility to

113 # call execute(sql, params, many, context).

114 self.execute_wrappers = []

115

116 self.client = self.client_class(self)

117 self.creation = self.creation_class(self)

118 self.features = self.features_class(self)

119 self.introspection = self.introspection_class(self)

120 self.ops = self.ops_class(self)

121 self.validation = self.validation_class(self)

122

123 def __repr__(self):

124 return (

125 f"<{self.__class__.__qualname__} "

126 f"vendor={self.vendor! r} alias={self.alias! r}>"

127 )

128

129 def ensure_timezone(self):

130 """

131 Ensure the connection's timezone is set to `self.timezone_name` and

132 return whether it changed or not.

133 """

134 return False

135

136 @cached_property

137 def timezone(self):

138 """

139 Return a tzinfo of the database connection time zone.

140

141 This is only used when time zone support is enabled. When a datetime is

142 read from the database, it is always returned in this time zone.

143

144 When the database backend supports time zones, it doesn't matter which

145 time zone Plain uses, as long as aware datetimes are used everywhere.

146 Other users connecting to the database can choose their own time zone.

147

148 When the database backend doesn't support time zones, the time zone

149 Plain uses may be constrained by the requirements of other users of

150 the database.

151 """

152 if self.settings_dict["TIME_ZONE"] is None:

153 return datetime.UTC

154 else:

155 return zoneinfo.ZoneInfo(self.settings_dict["TIME_ZONE"])

156

157 @cached_property

158 def timezone_name(self):

159 """

160 Name of the time zone of the database connection.

161 """

162 if self.settings_dict["TIME_ZONE"] is None:

163 return "UTC"

164 else:

165 return self.settings_dict["TIME_ZONE"]

166

167 @property

168 def queries_logged(self):

169 return self.force_debug_cursor or settings.DEBUG

170

171 @property

172 def queries(self):

173 if len(self.queries_log) == self.queries_log.maxlen:

174 warnings.warn(

175 f"Limit for query logging exceeded, only the last {self.queries_log.maxlen} queries "

176 "will be returned."

177 )

178 return list(self.queries_log)

179

180 def get_database_version(self):

181 """Return a tuple of the database's version."""

182 raise NotImplementedError(

183 "subclasses of BaseDatabaseWrapper may require a get_database_version() "

184 "method."

185 )

186

187 def check_database_version_supported(self):

188 """

189 Raise an error if the database version isn't supported by this

190 version of Plain.

191 """

192 if (

193 self.features.minimum_database_version is not None

194 and self.get_database_version() < self.features.minimum_database_version

195 ):

196 db_version = ".".join(map(str, self.get_database_version()))

197 min_db_version = ".".join(map(str, self.features.minimum_database_version))

198 raise NotSupportedError(

199 f"{self.display_name} {min_db_version} or later is required "

200 f"(found {db_version})."

201 )

202

203 # ##### Backend-specific methods for creating connections and cursors #####

204

205 def get_connection_params(self):

206 """Return a dict of parameters suitable for get_new_connection."""

207 raise NotImplementedError(

208 "subclasses of BaseDatabaseWrapper may require a get_connection_params() "

209 "method"

210 )

211

212 def get_new_connection(self, conn_params):

213 """Open a connection to the database."""

214 raise NotImplementedError(

215 "subclasses of BaseDatabaseWrapper may require a get_new_connection() "

216 "method"

217 )

218

219 def init_connection_state(self):

220 """Initialize the database connection settings."""

221 global RAN_DB_VERSION_CHECK

222 if self.alias not in RAN_DB_VERSION_CHECK:

223 self.check_database_version_supported()

224 RAN_DB_VERSION_CHECK.add(self.alias)

225

226 def create_cursor(self, name=None):

227 """Create a cursor. Assume that a connection is established."""

228 raise NotImplementedError(

229 "subclasses of BaseDatabaseWrapper may require a create_cursor() method"

230 )

231

232 # ##### Backend-specific methods for creating connections #####

233

234 def connect(self):

235 """Connect to the database. Assume that the connection is closed."""

236 # In case the previous connection was closed while in an atomic block

237 self.in_atomic_block = False

238 self.savepoint_ids = []

239 self.atomic_blocks = []

240 self.needs_rollback = False

241 # Reset parameters defining when to close/health-check the connection.

242 self.health_check_enabled = self.settings_dict["CONN_HEALTH_CHECKS"]

243 max_age = self.settings_dict["CONN_MAX_AGE"]

244 self.close_at = None if max_age is None else time.monotonic() + max_age

245 self.closed_in_transaction = False

246 self.errors_occurred = False

247 # New connections are healthy.

248 self.health_check_done = True

249 # Establish the connection

250 conn_params = self.get_connection_params()

251 self.connection = self.get_new_connection(conn_params)

252 self.set_autocommit(self.settings_dict["AUTOCOMMIT"])

253 self.init_connection_state()

254 connection_created.send(sender=self.__class__, connection=self)

255

256 self.run_on_commit = []

257

258 def ensure_connection(self):

259 """Guarantee that a connection to the database is established."""

260 if self.connection is None:

261 with self.wrap_database_errors:

262 self.connect()

263

264 # ##### Backend-specific wrappers for PEP-249 connection methods #####

265

266 def _prepare_cursor(self, cursor):

267 """

268 Validate the connection is usable and perform database cursor wrapping.

269 """

270 self.validate_thread_sharing()

271 if self.queries_logged:

272 wrapped_cursor = self.make_debug_cursor(cursor)

273 else:

274 wrapped_cursor = self.make_cursor(cursor)

275 return wrapped_cursor

276

277 def _cursor(self, name=None):

278 self.close_if_health_check_failed()

279 self.ensure_connection()

280 with self.wrap_database_errors:

281 return self._prepare_cursor(self.create_cursor(name))

282

283 def _commit(self):

284 if self.connection is not None:

285 with debug_transaction(self, "COMMIT"), self.wrap_database_errors:

286 return self.connection.commit()

287

288 def _rollback(self):

289 if self.connection is not None:

290 with debug_transaction(self, "ROLLBACK"), self.wrap_database_errors:

291 return self.connection.rollback()

292

293 def _close(self):

294 if self.connection is not None:

295 with self.wrap_database_errors:

296 return self.connection.close()

297

298 # ##### Generic wrappers for PEP-249 connection methods #####

299

300 def cursor(self):

301 """Create a cursor, opening a connection if necessary."""

302 return self._cursor()

303

304 def commit(self):

305 """Commit a transaction and reset the dirty flag."""

306 self.validate_thread_sharing()

307 self.validate_no_atomic_block()

308 self._commit()

309 # A successful commit means that the database connection works.

310 self.errors_occurred = False

311 self.run_commit_hooks_on_set_autocommit_on = True

312

313 def rollback(self):

314 """Roll back a transaction and reset the dirty flag."""

315 self.validate_thread_sharing()

316 self.validate_no_atomic_block()

317 self._rollback()

318 # A successful rollback means that the database connection works.

319 self.errors_occurred = False

320 self.needs_rollback = False

321 self.run_on_commit = []

322

323 def close(self):

324 """Close the connection to the database."""

325 self.validate_thread_sharing()

326 self.run_on_commit = []

327

328 # Don't call validate_no_atomic_block() to avoid making it difficult

329 # to get rid of a connection in an invalid state. The next connect()

330 # will reset the transaction state anyway.

331 if self.closed_in_transaction or self.connection is None:

332 return

333 try:

334 self._close()

335 finally:

336 if self.in_atomic_block:

337 self.closed_in_transaction = True

338 self.needs_rollback = True

339 else:

340 self.connection = None

341

342 # ##### Backend-specific savepoint management methods #####

343

344 def _savepoint(self, sid):

345 with self.cursor() as cursor:

346 cursor.execute(self.ops.savepoint_create_sql(sid))

347

348 def _savepoint_rollback(self, sid):

349 with self.cursor() as cursor:

350 cursor.execute(self.ops.savepoint_rollback_sql(sid))

351

352 def _savepoint_commit(self, sid):

353 with self.cursor() as cursor:

354 cursor.execute(self.ops.savepoint_commit_sql(sid))

355

356 def _savepoint_allowed(self):

357 # Savepoints cannot be created outside a transaction

358 return self.features.uses_savepoints and not self.get_autocommit()

359

360 # ##### Generic savepoint management methods #####

361

362 def savepoint(self):

363 """

364 Create a savepoint inside the current transaction. Return an

365 identifier for the savepoint that will be used for the subsequent

366 rollback or commit. Do nothing if savepoints are not supported.

367 """

368 if not self._savepoint_allowed():

369 return

370

371 thread_ident = _thread.get_ident()

372 tid = str(thread_ident).replace("-", "")

373

374 self.savepoint_state += 1

375 sid = "s%s_x%d" % (tid, self.savepoint_state)

376

377 self.validate_thread_sharing()

378 self._savepoint(sid)

379

380 return sid

381

382 def savepoint_rollback(self, sid):

383 """

384 Roll back to a savepoint. Do nothing if savepoints are not supported.

385 """

386 if not self._savepoint_allowed():

387 return

388

389 self.validate_thread_sharing()

390 self._savepoint_rollback(sid)

391

392 # Remove any callbacks registered while this savepoint was active.

393 self.run_on_commit = [

394 (sids, func, robust)

395 for (sids, func, robust) in self.run_on_commit

396 if sid not in sids

397 ]

398

399 def savepoint_commit(self, sid):

400 """

401 Release a savepoint. Do nothing if savepoints are not supported.

402 """

403 if not self._savepoint_allowed():

404 return

405

406 self.validate_thread_sharing()

407 self._savepoint_commit(sid)

408

409 def clean_savepoints(self):

410 """

411 Reset the counter used to generate unique savepoint ids in this thread.

412 """

413 self.savepoint_state = 0

414

415 # ##### Backend-specific transaction management methods #####

416

417 def _set_autocommit(self, autocommit):

418 """

419 Backend-specific implementation to enable or disable autocommit.

420 """

421 raise NotImplementedError(

422 "subclasses of BaseDatabaseWrapper may require a _set_autocommit() method"

423 )

424

425 # ##### Generic transaction management methods #####

426

427 def get_autocommit(self):

428 """Get the autocommit state."""

429 self.ensure_connection()

430 return self.autocommit

431

432 def set_autocommit(

433 self, autocommit, force_begin_transaction_with_broken_autocommit=False

434 ):

435 """

436 Enable or disable autocommit.

437

438 The usual way to start a transaction is to turn autocommit off.

439 SQLite does not properly start a transaction when disabling

440 autocommit. To avoid this buggy behavior and to actually enter a new

441 transaction, an explicit BEGIN is required. Using

442 force_begin_transaction_with_broken_autocommit=True will issue an

443 explicit BEGIN with SQLite. This option will be ignored for other

444 backends.

445 """

446 self.validate_no_atomic_block()

447 self.close_if_health_check_failed()

448 self.ensure_connection()

449

450 start_transaction_under_autocommit = (

451 force_begin_transaction_with_broken_autocommit

452 and not autocommit

453 and hasattr(self, "_start_transaction_under_autocommit")

454 )

455

456 if start_transaction_under_autocommit:

457 self._start_transaction_under_autocommit()

458 elif autocommit:

459 self._set_autocommit(autocommit)

460 else:

461 with debug_transaction(self, "BEGIN"):

462 self._set_autocommit(autocommit)

463 self.autocommit = autocommit

464

465 if autocommit and self.run_commit_hooks_on_set_autocommit_on:

466 self.run_and_clear_commit_hooks()

467 self.run_commit_hooks_on_set_autocommit_on = False

468

469 def get_rollback(self):

470 """Get the "needs rollback" flag -- for *advanced use* only."""

471 if not self.in_atomic_block:

472 raise TransactionManagementError(

473 "The rollback flag doesn't work outside of an 'atomic' block."

474 )

475 return self.needs_rollback

476

477 def set_rollback(self, rollback):

478 """

479 Set or unset the "needs rollback" flag -- for *advanced use* only.

480 """

481 if not self.in_atomic_block:

482 raise TransactionManagementError(

483 "The rollback flag doesn't work outside of an 'atomic' block."

484 )

485 self.needs_rollback = rollback

486

487 def validate_no_atomic_block(self):

488 """Raise an error if an atomic block is active."""

489 if self.in_atomic_block:

490 raise TransactionManagementError(

491 "This is forbidden when an 'atomic' block is active."

492 )

493

494 def validate_no_broken_transaction(self):

495 if self.needs_rollback:

496 raise TransactionManagementError(

497 "An error occurred in the current transaction. You can't "

498 "execute queries until the end of the 'atomic' block."

499 ) from self.rollback_exc

500

501 # ##### Foreign key constraints checks handling #####

502

503 @contextmanager

504 def constraint_checks_disabled(self):

505 """

506 Disable foreign key constraint checking.

507 """

508 disabled = self.disable_constraint_checking()

509 try:

510 yield

511 finally:

512 if disabled:

513 self.enable_constraint_checking()

514

515 def disable_constraint_checking(self):

516 """

517 Backends can implement as needed to temporarily disable foreign key

518 constraint checking. Should return True if the constraints were

519 disabled and will need to be reenabled.

520 """

521 return False

522

523 def enable_constraint_checking(self):

524 """

525 Backends can implement as needed to re-enable foreign key constraint

526 checking.

527 """

528 pass

529

530 def check_constraints(self, table_names=None):

531 """

532 Backends can override this method if they can apply constraint

533 checking (e.g. via "SET CONSTRAINTS ALL IMMEDIATE"). Should raise an

534 IntegrityError if any invalid foreign key references are encountered.

535 """

536 pass

537

538 # ##### Connection termination handling #####

539

540 def is_usable(self):

541 """

542 Test if the database connection is usable.

543

544 This method may assume that self.connection is not None.

545

546 Actual implementations should take care not to raise exceptions

547 as that may prevent Plain from recycling unusable connections.

548 """

549 raise NotImplementedError(

550 "subclasses of BaseDatabaseWrapper may require an is_usable() method"

551 )

552

553 def close_if_health_check_failed(self):

554 """Close existing connection if it fails a health check."""

555 if (

556 self.connection is None

557 or not self.health_check_enabled

558 or self.health_check_done

559 ):

560 return

561

562 if not self.is_usable():

563 self.close()

564 self.health_check_done = True

565

566 def close_if_unusable_or_obsolete(self):

567 """

568 Close the current connection if unrecoverable errors have occurred

569 or if it outlived its maximum age.

570 """

571 if self.connection is not None:

572 self.health_check_done = False

573 # If the application didn't restore the original autocommit setting,

574 # don't take chances, drop the connection.

575 if self.get_autocommit() != self.settings_dict["AUTOCOMMIT"]:

576 self.close()

577 return

578

579 # If an exception other than DataError or IntegrityError occurred

580 # since the last commit / rollback, check if the connection works.

581 if self.errors_occurred:

582 if self.is_usable():

583 self.errors_occurred = False

584 self.health_check_done = True

585 else:

586 self.close()

587 return

588

589 if self.close_at is not None and time.monotonic() >= self.close_at:

590 self.close()

591 return

592

593 # ##### Thread safety handling #####

594

595 @property

596 def allow_thread_sharing(self):

597 with self._thread_sharing_lock:

598 return self._thread_sharing_count > 0

599

600 def inc_thread_sharing(self):

601 with self._thread_sharing_lock:

602 self._thread_sharing_count += 1

603

604 def dec_thread_sharing(self):

605 with self._thread_sharing_lock:

606 if self._thread_sharing_count <= 0:

607 raise RuntimeError(

608 "Cannot decrement the thread sharing count below zero."

609 )

610 self._thread_sharing_count -= 1

611

612 def validate_thread_sharing(self):

613 """

614 Validate that the connection isn't accessed by another thread than the

615 one which originally created it, unless the connection was explicitly

616 authorized to be shared between threads (via the `inc_thread_sharing()`

617 method). Raise an exception if the validation fails.

618 """

619 if not (self.allow_thread_sharing or self._thread_ident == _thread.get_ident()):

620 raise DatabaseError(

621 "DatabaseWrapper objects created in a "

622 "thread can only be used in that same thread. The object "

623 f"with alias '{self.alias}' was created in thread id {self._thread_ident} and this is "

624 f"thread id {_thread.get_ident()}."

625 )

626

627 # ##### Miscellaneous #####

628

629 def prepare_database(self):

630 """

631 Hook to do any database check or preparation, generally called before

632 migrating a project or an app.

633 """

634 pass

635

636 @cached_property

637 def wrap_database_errors(self):

638 """

639 Context manager and decorator that re-throws backend-specific database

640 exceptions using Plain's common wrappers.

641 """

642 return DatabaseErrorWrapper(self)

643

644 def chunked_cursor(self):

645 """

646 Return a cursor that tries to avoid caching in the database (if

647 supported by the database), otherwise return a regular cursor.

648 """

649 return self.cursor()

650

651 def make_debug_cursor(self, cursor):

652 """Create a cursor that logs all queries in self.queries_log."""

653 return utils.CursorDebugWrapper(cursor, self)

654

655 def make_cursor(self, cursor):

656 """Create a cursor without debug logging."""

657 return utils.CursorWrapper(cursor, self)

658

659 @contextmanager

660 def temporary_connection(self):

661 """

662 Context manager that ensures that a connection is established, and

663 if it opened one, closes it to avoid leaving a dangling connection.

664 This is useful for operations outside of the request-response cycle.

665

666 Provide a cursor: with self.temporary_connection() as cursor: ...

667 """

668 must_close = self.connection is None

669 try:

670 with self.cursor() as cursor:

671 yield cursor

672 finally:

673 if must_close:

674 self.close()

675

676 @contextmanager

677 def _nodb_cursor(self):

678 """

679 Return a cursor from an alternative connection to be used when there is

680 no need to access the main database, specifically for test db

681 creation/deletion. This also prevents the production database from

682 being exposed to potential child threads while (or after) the test

683 database is destroyed. Refs #10868, #17786, #16969.

684 """

685 conn = self.__class__({**self.settings_dict, "NAME": None}, alias=NO_DB_ALIAS)

686 try:

687 with conn.cursor() as cursor:

688 yield cursor

689 finally:

690 conn.close()

691

692 def schema_editor(self, *args, **kwargs):

693 """

694 Return a new instance of this backend's SchemaEditor.

695 """

696 if self.SchemaEditorClass is None:

697 raise NotImplementedError(

698 "The SchemaEditorClass attribute of this database wrapper is still None"

699 )

700 return self.SchemaEditorClass(self, *args, **kwargs)

701

702 def on_commit(self, func, robust=False):

703 if not callable(func):

704 raise TypeError("on_commit()'s callback must be a callable.")

705 if self.in_atomic_block:

706 # Transaction in progress; save for execution on commit.

707 self.run_on_commit.append((set(self.savepoint_ids), func, robust))

708 elif not self.get_autocommit():

709 raise TransactionManagementError(

710 "on_commit() cannot be used in manual transaction management"

711 )

712 else:

713 # No transaction in progress and in autocommit mode; execute

714 # immediately.

715 if robust:

716 try:

717 func()

718 except Exception as e:

719 logger.error(

720 f"Error calling {func.__qualname__} in on_commit() (%s).",

721 e,

722 exc_info=True,

723 )

724 else:

725 func()

726

727 def run_and_clear_commit_hooks(self):

728 self.validate_no_atomic_block()

729 current_run_on_commit = self.run_on_commit

730 self.run_on_commit = []

731 while current_run_on_commit:

732 _, func, robust = current_run_on_commit.pop(0)

733 if robust:

734 try:

735 func()

736 except Exception as e:

737 logger.error(

738 f"Error calling {func.__qualname__} in on_commit() during "

739 f"transaction (%s).",

740 e,

741 exc_info=True,

742 )

743 else:

744 func()

745

746 @contextmanager

747 def execute_wrapper(self, wrapper):

748 """

749 Return a context manager under which the wrapper is applied to suitable

750 database query executions.

751 """

752 self.execute_wrappers.append(wrapper)

753 try:

754 yield

755 finally:

756 self.execute_wrappers.pop()

757

758 def copy(self, alias=None):

759 """

760 Return a copy of this connection.

761

762 For tests that require two connections to the same database.

763 """

764 settings_dict = copy.deepcopy(self.settings_dict)

765 if alias is None:

766 alias = self.alias

767 return type(self)(settings_dict, alias)