Coverage for mongodb_odm / connection.py: 100%

1import logging 6 ctx1fabcde

2from typing import Any, Optional, Union, cast 6 ctx1fabcde

3 

4from mongodb_odm.exceptions import ConnectionError, InvalidAction, InvalidConnection 6 ctx1fabcde

5from mongodb_odm.types import DICT_TYPE 6 ctx1fabcde

6from mongodb_odm.utils._internal_models import Connection 6 ctx1fabcde

7from pymongo import AsyncMongoClient, MongoClient 6 ctx1fabcde

8from pymongo.asynchronous.database import AsyncDatabase 6 ctx1fabcde

9from pymongo.database import Database 6 ctx1fabcde

10 

11logger = logging.getLogger(__name__) 6 ctx1fabcde

12 

13"""Store database connection related values in this variable""" 5 ctx1abcde

14__connection_obj = Connection() 6 ctx1fabcde

15 

16 

17def _get_connection_client(url: str, connection_kwargs: DICT_TYPE) -> MongoClient[Any]: 6 ctx1fabcde

18 return MongoClient(url, **connection_kwargs) 6 ctx1fabcde

19 

20 

21def _get_async_connection_client( 6 ctx1fabcde

22 url: str, connection_kwargs: DICT_TYPE 

23) -> AsyncMongoClient[Any]: 

24 return AsyncMongoClient(url, **connection_kwargs) 6 ctx1fabcde

25 

26 

27def connect( 6 ctx1fabcde

28 url: str, 

29 databases: Optional[set[str]] = None, 

30 connection_kwargs: Optional[DICT_TYPE] = None, 

31 async_is_enabled: bool = False, 

32) -> Union[AsyncMongoClient[Any], MongoClient[Any]]: 

33 """ 

34 This connect function should manage and store database connection config 

35 that are passed by user. 

36 

37 url: string type, required 

38 This should be a valid mongodb connection string with default database. 

39 

40 databases: set of strings, optional 

41 values should be a set of strings that are assigned on define models. 

42 The value is meant to validate that user is not trying to connect 

43 to a database that is not defined in the models. 

44 

45 connection_kwargs: dict, optional 

46 Additional connection parameters that can be passed to MongoClient or AsyncMongoClient. 

47 

48 async_is_enabled: bool, optional 

49 If True then use AsyncMongoClient, otherwise use MongoClient. 

50 Default is False. 

51 """ 

52 global __connection_obj 

53 

54 if __connection_obj.client is not None: 6 ctx1fabcde

55 """Log a warning if a user tries to connect multiple times.""" 5 ctx1abcde

56 logger.warning("Already have an connection.") 6 ctx1fabcde

57 return __connection_obj.client 6 ctx1fabcde

58 

59 __connection_obj = Connection() 6 ctx1fabcde

60 

61 if databases is None: 6 ctx1fabcde

62 """Assign empty set as default value""" 5 ctx1abcde

63 databases = set() 6 ctx1fabcde

64 

65 if not connection_kwargs: 6 ctx1fabcde

66 connection_kwargs = {} 6 ctx1fabcde

67 

68 client: Union[AsyncMongoClient[Any], MongoClient[Any]] 

69 if async_is_enabled: 6 ctx1fabcde

70 client = _get_async_connection_client(url, connection_kwargs) 6 ctx1fabcde

71 else: 

72 client = _get_connection_client(url, connection_kwargs) 6 ctx1fabcde

73 

74 default_database = client.get_default_database().name 6 ctx1fabcde

75 databases.add(default_database) 6 ctx1fabcde

76 

77 """Store the user configuration in the global variable to use later.""" 5 ctx1abcde

78 __connection_obj.url = url 6 ctx1fabcde

79 __connection_obj.databases = databases 6 ctx1fabcde

80 __connection_obj.client = client 6 ctx1fabcde

81 __connection_obj.async_is_enabled = async_is_enabled 6 ctx1fabcde

82 __connection_obj.connection_kwargs = connection_kwargs 6 ctx1fabcde

83 

84 logger.info("Connection established successfully") 6 ctx1fabcde

85 

86 return __connection_obj.client 6 ctx1fabcde

87 

88 

89def _disconnect_common() -> bool: 6 ctx1fabcde

90 global __connection_obj 

91 

92 from mongodb_odm.models import _clear_cache 6 ctx1fabcde

93 

94 _clear_cache() 6 ctx1fabcde

95 

96 __connection_obj = Connection() # Reset the connection object 6 ctx1fabcde

97 

98 logger.info("Disconnect the db connection") 6 ctx1fabcde

99 

100 return True 6 ctx1fabcde

101 

102 

103def disconnect(raise_error: bool = True) -> bool: 6 ctx1fabcde

104 global __connection_obj 

105 

106 if __connection_obj.client is None: 6 ctx1fabcde

107 logger.warning("No client connection found") 6 ctx1fabcde

108 return _disconnect_common() 6 ctx1fabcde

109 

110 if isinstance(__connection_obj.client, AsyncMongoClient): 6 ctx1fabcde

111 if not raise_error: 6 ctx1fabcde

112 logger.warning( 6 ctx1fabcde

113 "The client is configured as async. Use adisconnect() instead." 

114 ) 

115 # We are not closing the async client expecting since this is a silent disconnect. 

116 # Also it's not possible to call close() on an async client in a sync context. 

117 __connection_obj.client = None 6 ctx1fabcde

118 

119 return _disconnect_common() 6 ctx1fabcde

120 else: 

121 raise InvalidAction( 6 ctx1fabcde

122 "The client is configured as async. Use adisconnect() instead." 

123 ) 

124 

125 __connection_obj.client.close() 6 ctx1fabcde

126 __connection_obj.client = None 6 ctx1fabcde

127 

128 return _disconnect_common() 6 ctx1fabcde

129 

130 

131async def adisconnect(raise_error: bool = True) -> bool: 6 ctx1fabcde

132 global __connection_obj 

133 

134 if __connection_obj.client is None: 6 ctx1fabcde

135 logger.warning("No client connection found") 6 ctx1fabcde

136 return _disconnect_common() 6 ctx1fabcde

137 

138 if isinstance(__connection_obj.client, MongoClient): 6 ctx1fabcde

139 if not raise_error: 6 ctx1fabcde

140 return disconnect() 6 ctx1fabcde

141 else: 

142 raise InvalidAction( 6 ctx1fabcde

143 "The client is configured as sync. Use disconnect() instead." 

144 ) 

145 

146 await __connection_obj.client.close() 6 ctx1fabcde

147 __connection_obj.client = None 6 ctx1fabcde

148 

149 return _disconnect_common() 6 ctx1fabcde

150 

151 

152def get_client() -> Union[AsyncMongoClient[Any], MongoClient[Any]]: 6 ctx1fabcde

153 """ 

154 Function should return MongoClient if url exists. 

155 Otherwise raise ConnectionError error 

156 """ 

157 global __connection_obj 

158 

159 if __connection_obj.client: 6 ctx1fabcde

160 return __connection_obj.client 6 ctx1fabcde

161 

162 if __connection_obj.url is None: 6 ctx1fabcde

163 raise ConnectionError("DB connection URL is not provided") 6 ctx1fabcde

164 

165 databases = __connection_obj.databases 6 ctx1fabcde

166 connection_kwargs = __connection_obj.connection_kwargs 6 ctx1fabcde

167 async_is_enabled = __connection_obj.async_is_enabled 6 ctx1fabcde

168 

169 return connect( 6 ctx1fabcde

170 __connection_obj.url, 

171 databases=databases, 

172 connection_kwargs=connection_kwargs, 

173 async_is_enabled=async_is_enabled, 

174 ) 

175 

176 

177def db( 6 ctx1fabcde

178 database: Optional[str] = None, 

179 is_async_action: Optional[bool] = None, 

180) -> Union[Database[Any], AsyncDatabase[Any]]: 

181 global __connection_obj 

182 

183 if is_async_action is None: 6 ctx1fabcde

184 is_async_action = __connection_obj.async_is_enabled 6 ctx1fabcde

185 

186 client = get_client() 6 ctx1fabcde

187 

188 if isinstance(client, MongoClient) and is_async_action is True: 6 ctx1fabcde

189 raise InvalidAction("The client is not configured as async") 6 ctx1fabcde

190 elif isinstance(client, AsyncMongoClient) and is_async_action is False: 6 ctx1fabcde

191 raise InvalidAction("The client is not configured as sync") 6 ctx1fabcde

192 

193 if database: 6 ctx1fabcde

194 if not __connection_obj.databases or database not in __connection_obj.databases: 6 ctx1fabcde

195 """ 5 ctx1abcde

196 Make sure valid database strings are passed 

197 and avoid creating invalid databases by mistake 

198 because mongodb creates databases on the fly. 

199 """ 

200 raise InvalidConnection(f'Invalid database key was passed "{database}"') 6 ctx1fabcde

201 

202 return client[database] 6 ctx1fabcde

203 else: 

204 """ 5 ctx1abcde

205 If database string was not passed then return default database 

206 that is provided in connection url. 

207 """ 

208 return client.get_database() 6 ctx1fabcde

209 

210 

211def drop_database(database: Optional[str] = None) -> None: 6 ctx1fabcde

212 collection = db(database, is_async_action=False) 6 ctx1fabcde

213 

214 collection.command("dropDatabase") 6 ctx1fabcde

215 

216 

217async def adrop_database(database: Optional[str] = None) -> None: # noqa 6 ctx1fabcde

218 collection = db(database, is_async_action=True) 6 ctx1fabcde

219 

220 collection = cast(AsyncDatabase[Any], collection) 6 ctx1fabcde

221 

222 await collection.command("dropDatabase") 6 ctx1fabcde

223 

224 

225def is_async() -> bool: 6 ctx1fabcde

226 global __connection_obj 

227 

228 if not __connection_obj.client: 6 ctx1fabcde

229 raise ConnectionError("No connection established. Please connect first.") 6 ctx1fabcde

230 

231 return isinstance(__connection_obj.client, AsyncMongoClient) 6 ctx1fabcde