Coverage for /Users/davegaeddert/Development/dropseed/plain/plain-models/plain/models/database_url.py: 36%

67 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-16 22:04 -0500

1# Copyright (c) Kenneth Reitz & individual contributors 

2# All rights reserved. 

3 

4# Redistribution and use in source and binary forms, with or without modification, 

5# are permitted provided that the following conditions are met: 

6 

7# 1. Redistributions of source code must retain the above copyright notice, 

8# this list of conditions and the following disclaimer. 

9 

10# 2. Redistributions in binary form must reproduce the above copyright 

11# notice, this list of conditions and the following disclaimer in the 

12# documentation and/or other materials provided with the distribution. 

13 

14# 3. Neither the name of Plain nor the names of its contributors may be used 

15# to endorse or promote products derived from this software without 

16# specific prior written permission. 

17 

18# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 

19# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 

20# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 

21# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR 

22# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 

23# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 

24# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 

25# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 

26# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 

27# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 

28import logging 

29import os 

30import urllib.parse as urlparse 

31from typing import Any, TypedDict 

32 

33DEFAULT_ENV = "DATABASE_URL" 

34 

35SCHEMES = { 

36 "postgres": "plain.models.backends.postgresql", 

37 "postgresql": "plain.models.backends.postgresql", 

38 "pgsql": "plain.models.backends.postgresql", 

39 "mysql": "plain.models.backends.mysql", 

40 "mysql2": "plain.models.backends.mysql", 

41 "sqlite": "plain.models.backends.sqlite3", 

42} 

43 

44# Register database schemes in URLs. 

45for key in SCHEMES.keys(): 

46 urlparse.uses_netloc.append(key) 

47 

48 

49class DBConfig(TypedDict, total=False): 

50 AUTOCOMMIT: bool 

51 CONN_MAX_AGE: int | None 

52 CONN_HEALTH_CHECKS: bool 

53 DISABLE_SERVER_SIDE_CURSORS: bool 

54 ENGINE: str 

55 HOST: str 

56 NAME: str 

57 OPTIONS: dict[str, Any] | None 

58 PASSWORD: str 

59 PORT: str | int 

60 TEST: dict[str, Any] 

61 TIME_ZONE: str 

62 USER: str 

63 

64 

65def config( 

66 env: str = DEFAULT_ENV, 

67 default: str | None = None, 

68 engine: str | None = None, 

69 conn_max_age: int | None = 0, 

70 conn_health_checks: bool = False, 

71 ssl_require: bool = False, 

72 test_options: dict | None = None, 

73) -> DBConfig: 

74 """Returns configured DATABASE dictionary from DATABASE_URL.""" 

75 s = os.environ.get(env, default) 

76 

77 if s is None: 

78 logging.warning( 

79 "No %s environment variable set, and so no databases setup" % env 

80 ) 

81 

82 if s: 

83 return parse( 

84 s, engine, conn_max_age, conn_health_checks, ssl_require, test_options 

85 ) 

86 

87 return {} 

88 

89 

90def parse( 

91 url: str, 

92 engine: str | None = None, 

93 conn_max_age: int | None = 0, 

94 conn_health_checks: bool = False, 

95 ssl_require: bool = False, 

96 test_options: dict | None = None, 

97) -> DBConfig: 

98 """Parses a database URL.""" 

99 if url == "sqlite://:memory:": 

100 # this is a special case, because if we pass this URL into 

101 # urlparse, urlparse will choke trying to interpret "memory" 

102 # as a port number 

103 return {"ENGINE": SCHEMES["sqlite"], "NAME": ":memory:"} 

104 # note: no other settings are required for sqlite 

105 

106 # otherwise parse the url as normal 

107 parsed_config: DBConfig = {} 

108 

109 if test_options is None: 

110 test_options = {} 

111 

112 spliturl = urlparse.urlsplit(url) 

113 

114 # Split query strings from path. 

115 path = spliturl.path[1:] 

116 query = urlparse.parse_qs(spliturl.query) 

117 

118 # If we are using sqlite and we have no path, then assume we 

119 # want an in-memory database (this is the behaviour of sqlalchemy) 

120 if spliturl.scheme == "sqlite" and path == "": 

121 path = ":memory:" 

122 

123 # Handle postgres percent-encoded paths. 

124 hostname = spliturl.hostname or "" 

125 if "%" in hostname: 

126 # Switch to url.netloc to avoid lower cased paths 

127 hostname = spliturl.netloc 

128 if "@" in hostname: 

129 hostname = hostname.rsplit("@", 1)[1] 

130 # Use URL Parse library to decode % encodes 

131 hostname = urlparse.unquote(hostname) 

132 

133 # Lookup specified engine. 

134 if engine is None: 

135 engine = SCHEMES.get(spliturl.scheme) 

136 if engine is None: 

137 raise ValueError( 

138 "No support for '{}'. We support: {}".format( 

139 spliturl.scheme, ", ".join(sorted(SCHEMES.keys())) 

140 ) 

141 ) 

142 

143 port = spliturl.port 

144 

145 # Update with environment configuration. 

146 parsed_config.update( 

147 { 

148 "NAME": urlparse.unquote(path or ""), 

149 "USER": urlparse.unquote(spliturl.username or ""), 

150 "PASSWORD": urlparse.unquote(spliturl.password or ""), 

151 "HOST": hostname, 

152 "PORT": port or "", 

153 "CONN_MAX_AGE": conn_max_age, 

154 "CONN_HEALTH_CHECKS": conn_health_checks, 

155 "ENGINE": engine, 

156 } 

157 ) 

158 if test_options: 

159 parsed_config.update( 

160 { 

161 "TEST": test_options, 

162 } 

163 ) 

164 

165 # Pass the query string into OPTIONS. 

166 options: dict[str, Any] = {} 

167 for key, values in query.items(): 

168 if spliturl.scheme == "mysql" and key == "ssl-ca": 

169 options["ssl"] = {"ca": values[-1]} 

170 continue 

171 

172 options[key] = values[-1] 

173 

174 if ssl_require: 

175 options["sslmode"] = "require" 

176 

177 # Support for Postgres Schema URLs 

178 if "currentSchema" in options and engine in ( 

179 "plain.models.backends.postgresql_psycopg2", 

180 "plain.models.backends.postgresql", 

181 ): 

182 options["options"] = "-c search_path={}".format(options.pop("currentSchema")) 

183 

184 if options: 

185 parsed_config["OPTIONS"] = options 

186 

187 return parsed_config