Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/backends/sqlite3/creation.py: 33%
98 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 11:16 -0600
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 11:16 -0600
1import multiprocessing
2import os
3import shutil
4import sqlite3
5import sys
6from pathlib import Path
8from plain.models.backends.base.creation import BaseDatabaseCreation
9from plain.models.db import NotSupportedError
12class DatabaseCreation(BaseDatabaseCreation):
13 @staticmethod
14 def is_in_memory_db(database_name):
15 return not isinstance(database_name, Path) and (
16 database_name == ":memory:" or "mode=memory" in database_name
17 )
19 def _get_test_db_name(self):
20 test_database_name = self.connection.settings_dict["TEST"]["NAME"] or ":memory:"
21 if test_database_name == ":memory:":
22 return f"file:memorydb_{self.connection.alias}?mode=memory&cache=shared"
23 return test_database_name
25 def _create_test_db(self, verbosity, autoclobber, keepdb=False):
26 test_database_name = self._get_test_db_name()
28 if keepdb:
29 return test_database_name
30 if not self.is_in_memory_db(test_database_name):
31 # Erase the old test database
32 if verbosity >= 1:
33 self.log(
34 f"Destroying old test database for alias {self._get_database_display_str(verbosity, test_database_name)}..."
35 )
36 if os.access(test_database_name, os.F_OK):
37 if not autoclobber:
38 confirm = input(
39 "Type 'yes' if you would like to try deleting the test "
40 f"database '{test_database_name}', or 'no' to cancel: "
41 )
42 if autoclobber or confirm == "yes":
43 try:
44 os.remove(test_database_name)
45 except Exception as e:
46 self.log(f"Got an error deleting the old test database: {e}")
47 sys.exit(2)
48 else:
49 self.log("Tests cancelled.")
50 sys.exit(1)
51 return test_database_name
53 def get_test_db_clone_settings(self, suffix):
54 orig_settings_dict = self.connection.settings_dict
55 source_database_name = orig_settings_dict["NAME"]
57 if not self.is_in_memory_db(source_database_name):
58 root, ext = os.path.splitext(source_database_name)
59 return {**orig_settings_dict, "NAME": f"{root}_{suffix}{ext}"}
61 start_method = multiprocessing.get_start_method()
62 if start_method == "fork":
63 return orig_settings_dict
64 if start_method == "spawn":
65 return {
66 **orig_settings_dict,
67 "NAME": f"{self.connection.alias}_{suffix}.sqlite3",
68 }
69 raise NotSupportedError(
70 f"Cloning with start method {start_method!r} is not supported."
71 )
73 def _clone_test_db(self, suffix, verbosity, keepdb=False):
74 source_database_name = self.connection.settings_dict["NAME"]
75 target_database_name = self.get_test_db_clone_settings(suffix)["NAME"]
76 if not self.is_in_memory_db(source_database_name):
77 # Erase the old test database
78 if os.access(target_database_name, os.F_OK):
79 if keepdb:
80 return
81 if verbosity >= 1:
82 self.log(
83 "Destroying old test database for alias {}...".format(
84 self._get_database_display_str(
85 verbosity, target_database_name
86 ),
87 )
88 )
89 try:
90 os.remove(target_database_name)
91 except Exception as e:
92 self.log(f"Got an error deleting the old test database: {e}")
93 sys.exit(2)
94 try:
95 shutil.copy(source_database_name, target_database_name)
96 except Exception as e:
97 self.log(f"Got an error cloning the test database: {e}")
98 sys.exit(2)
99 # Forking automatically makes a copy of an in-memory database.
100 # Spawn requires migrating to disk which will be re-opened in
101 # setup_worker_connection.
102 elif multiprocessing.get_start_method() == "spawn":
103 ondisk_db = sqlite3.connect(target_database_name, uri=True)
104 self.connection.connection.backup(ondisk_db)
105 ondisk_db.close()
107 def _destroy_test_db(self, test_database_name, verbosity):
108 if test_database_name and not self.is_in_memory_db(test_database_name):
109 # Remove the SQLite database file
110 os.remove(test_database_name)
112 def test_db_signature(self):
113 """
114 Return a tuple that uniquely identifies a test database.
116 This takes into account the special cases of ":memory:" and "" for
117 SQLite since the databases will be distinct despite having the same
118 TEST NAME. See https://www.sqlite.org/inmemorydb.html
119 """
120 test_database_name = self._get_test_db_name()
121 sig = [self.connection.settings_dict["NAME"]]
122 if self.is_in_memory_db(test_database_name):
123 sig.append(self.connection.alias)
124 else:
125 sig.append(test_database_name)
126 return tuple(sig)
128 def setup_worker_connection(self, _worker_id):
129 settings_dict = self.get_test_db_clone_settings(_worker_id)
130 # connection.settings_dict must be updated in place for changes to be
131 # reflected in plain.models.connections. Otherwise new threads would
132 # connect to the default database instead of the appropriate clone.
133 start_method = multiprocessing.get_start_method()
134 if start_method == "fork":
135 # Update settings_dict in place.
136 self.connection.settings_dict.update(settings_dict)
137 self.connection.close()
138 elif start_method == "spawn":
139 alias = self.connection.alias
140 connection_str = (
141 f"file:memorydb_{alias}_{_worker_id}?mode=memory&cache=shared"
142 )
143 source_db = self.connection.Database.connect(
144 f"file:{alias}_{_worker_id}.sqlite3", uri=True
145 )
146 target_db = sqlite3.connect(connection_str, uri=True)
147 source_db.backup(target_db)
148 source_db.close()
149 # Update settings_dict in place.
150 self.connection.settings_dict.update(settings_dict)
151 self.connection.settings_dict["NAME"] = connection_str
152 # Re-open connection to in-memory database before closing copy
153 # connection.
154 self.connection.connect()
155 target_db.close()