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

67 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-23 11:16 -0600

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(f"No {env} environment variable set, and so no databases setup") 

79 

80 if s: 

81 return parse( 

82 s, engine, conn_max_age, conn_health_checks, ssl_require, test_options 

83 ) 

84 

85 return {} 

86 

87 

88def parse( 

89 url: str, 

90 engine: str | None = None, 

91 conn_max_age: int | None = 0, 

92 conn_health_checks: bool = False, 

93 ssl_require: bool = False, 

94 test_options: dict | None = None, 

95) -> DBConfig: 

96 """Parses a database URL.""" 

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

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

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

100 # as a port number 

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

102 # note: no other settings are required for sqlite 

103 

104 # otherwise parse the url as normal 

105 parsed_config: DBConfig = {} 

106 

107 if test_options is None: 

108 test_options = {} 

109 

110 spliturl = urlparse.urlsplit(url) 

111 

112 # Split query strings from path. 

113 path = spliturl.path[1:] 

114 query = urlparse.parse_qs(spliturl.query) 

115 

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

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

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

119 path = ":memory:" 

120 

121 # Handle postgres percent-encoded paths. 

122 hostname = spliturl.hostname or "" 

123 if "%" in hostname: 

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

125 hostname = spliturl.netloc 

126 if "@" in hostname: 

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

128 # Use URL Parse library to decode % encodes 

129 hostname = urlparse.unquote(hostname) 

130 

131 # Lookup specified engine. 

132 if engine is None: 

133 engine = SCHEMES.get(spliturl.scheme) 

134 if engine is None: 

135 raise ValueError( 

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

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

138 ) 

139 ) 

140 

141 port = spliturl.port 

142 

143 # Update with environment configuration. 

144 parsed_config.update( 

145 { 

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

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

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

149 "HOST": hostname, 

150 "PORT": port or "", 

151 "CONN_MAX_AGE": conn_max_age, 

152 "CONN_HEALTH_CHECKS": conn_health_checks, 

153 "ENGINE": engine, 

154 } 

155 ) 

156 if test_options: 

157 parsed_config.update( 

158 { 

159 "TEST": test_options, 

160 } 

161 ) 

162 

163 # Pass the query string into OPTIONS. 

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

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

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

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

168 continue 

169 

170 options[key] = values[-1] 

171 

172 if ssl_require: 

173 options["sslmode"] = "require" 

174 

175 # Support for Postgres Schema URLs 

176 if "currentSchema" in options and engine in ( 

177 "plain.models.backends.postgresql_psycopg2", 

178 "plain.models.backends.postgresql", 

179 ): 

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

181 

182 if options: 

183 parsed_config["OPTIONS"] = options 

184 

185 return parsed_config