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

1import multiprocessing 

2import os 

3import shutil 

4import sqlite3 

5import sys 

6from pathlib import Path 

7 

8from plain.models.backends.base.creation import BaseDatabaseCreation 

9from plain.models.db import NotSupportedError 

10 

11 

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 ) 

18 

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 

24 

25 def _create_test_db(self, verbosity, autoclobber, keepdb=False): 

26 test_database_name = self._get_test_db_name() 

27 

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 

52 

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"] 

56 

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}"} 

60 

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 ) 

72 

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() 

106 

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) 

111 

112 def test_db_signature(self): 

113 """ 

114 Return a tuple that uniquely identifies a test database. 

115 

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) 

127 

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()