Coverage for src/prisma/engine/utils.py: 88%

112 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-04-28 15:17 +0000

1from __future__ import annotations 

2 

3import os 

4import sys 

5import time 

6import socket 

7import logging 

8import subprocess 

9from typing import Any, Dict, Type, NoReturn 

10from pathlib import Path 

11 

12from . import errors 

13from .. import config, errors as prisma_errors 

14from ..utils import DEBUG_GENERATOR, time_since 

15from ..binaries import platform 

16from ..http_abstract import AbstractResponse 

17 

18log: logging.Logger = logging.getLogger(__name__) 

19 

20ERROR_MAPPING: Dict[str, Type[Exception]] = { 

21 'P2002': prisma_errors.UniqueViolationError, 

22 'P2003': prisma_errors.ForeignKeyViolationError, 

23 'P2009': prisma_errors.FieldNotFoundError, 

24 'P2010': prisma_errors.RawQueryError, 

25 'P2012': prisma_errors.MissingRequiredValueError, 

26 'P2019': prisma_errors.InputError, 

27 'P2021': prisma_errors.TableNotFoundError, 

28 'P2025': prisma_errors.RecordNotFoundError, 

29} 

30 

31META_ERROR_MAPPING: dict[str, type[Exception]] = { 

32 'UnknownArgument': prisma_errors.FieldNotFoundError, 

33 'UnknownInputField': prisma_errors.FieldNotFoundError, 

34 'UnknownSelectionField': prisma_errors.FieldNotFoundError, 

35} 

36 

37 

38def query_engine_name() -> str: 

39 return f'prisma-query-engine-{platform.check_for_extension(platform.binary_platform())}' 

40 

41 

42def _resolve_from_binary_paths(binary_paths: dict[str, str]) -> Path | None: 

43 if config.binary_platform is not None: 43 ↛ 44line 43 didn't jump to line 44, because the condition on line 43 was never true

44 return Path(binary_paths[config.binary_platform]) 

45 

46 paths = [Path(p) for p in binary_paths.values()] 

47 

48 # fast path for when there are no `binaryTargets` defined 

49 if len(paths) == 1: 

50 return paths[0] 

51 

52 for path in paths: 52 ↛ 56line 52 didn't jump to line 56, because the loop on line 52 never started

53 # we only want to resolve to binary's that we can run on the current machine. 

54 # because of the `binaryTargets` option some of the binaries we are given may 

55 # not be targeting the same architecture as the current machine 

56 if path.exists() and _can_execute_binary(path): 

57 return path 

58 

59 # none of the given paths existed or they target a different architecture 

60 return None 

61 

62 

63def _can_execute_binary(path: Path) -> bool: 

64 proc = subprocess.run([str(path), '--version'], check=False) 

65 log.debug('Executable check for %s exited with code: %s', path, proc.returncode) 

66 return proc.returncode == 0 

67 

68 

69def ensure(binary_paths: dict[str, str]) -> Path: 

70 start_time = time.monotonic() 

71 file = None 

72 force_version = not DEBUG_GENERATOR 

73 name = query_engine_name() 

74 local_path = Path.cwd().joinpath(name) 

75 global_path = config.binary_cache_dir.joinpath(name) 

76 file_from_paths = _resolve_from_binary_paths(binary_paths) 

77 

78 log.debug('Expecting local query engine %s', local_path) 

79 log.debug('Expecting global query engine %s', global_path) 

80 

81 binary = os.environ.get('PRISMA_QUERY_ENGINE_BINARY') 

82 if binary: 

83 log.debug('PRISMA_QUERY_ENGINE_BINARY is defined, using %s', binary) 

84 

85 if not Path(binary).exists(): 

86 raise errors.BinaryNotFoundError( 

87 'PRISMA_QUERY_ENGINE_BINARY was provided, ' f'but no query engine was found at {binary}' 

88 ) 

89 

90 file = Path(binary) 

91 force_version = False 

92 elif local_path.exists(): 

93 file = local_path 

94 log.debug('Query engine found in the working directory') 

95 elif file_from_paths is not None and file_from_paths.exists(): 

96 file = file_from_paths 

97 log.debug( 

98 'Query engine found from the Prisma CLI generated path: %s', 

99 file_from_paths, 

100 ) 

101 elif global_path.exists(): 101 ↛ 102line 101 didn't jump to line 102, because the condition on line 101 was never true

102 file = global_path 

103 log.debug('Query engine found in the global path') 

104 

105 if not file: 

106 if file_from_paths is not None: 

107 expected = f'{local_path}, {global_path} or {file_from_paths} to exist but none' 

108 else: 

109 expected = f'{local_path} or {global_path} to exist but neither' 

110 

111 raise errors.BinaryNotFoundError( 

112 f'Expected {expected} were found or could not be executed.\n' + 'Try running prisma py fetch' 

113 ) 

114 

115 log.debug('Using Query Engine binary at %s', file) 

116 

117 start_version = time.monotonic() 

118 process = subprocess.run([str(file.absolute()), '--version'], stdout=subprocess.PIPE, check=True) 

119 log.debug('Version check took %s', time_since(start_version)) 

120 

121 version = str(process.stdout, sys.getdefaultencoding()).replace('query-engine', '').strip() 

122 log.debug('Using query engine version %s', version) 

123 

124 if force_version and version != config.expected_engine_version: 

125 raise errors.MismatchedVersionsError(expected=config.expected_engine_version, got=version) 

126 

127 log.debug('Using query engine at %s', file) 

128 log.debug('Ensuring query engine took: %s', time_since(start_time)) 

129 return file 

130 

131 

132def get_open_port() -> int: 

133 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 

134 sock.bind(('', 0)) 

135 port = sock.getsockname()[1] 

136 sock.close() 

137 return int(port) 

138 

139 

140def handle_response_errors(resp: AbstractResponse[Any], data: Any) -> NoReturn: 

141 for error in data: 

142 try: 

143 base_error_message = error.get('error', '') 

144 user_facing = error.get('user_facing_error', {}) 

145 code = user_facing.get('error_code') 

146 if code is None: 146 ↛ 147line 146 didn't jump to line 147, because the condition on line 146 was never true

147 continue 

148 

149 # TODO: the order of these if statements is important because 

150 # the P2009 code can be returned for both: missing a required value 

151 # and an unknown field error. As we want to explicitly handle 

152 # the missing a required value error then we need to check for that first. 

153 # As we can only check for this error by searching the message then this 

154 # comes with performance concerns. 

155 message = user_facing.get('message', '') 

156 

157 if code == 'P2028': 

158 if base_error_message.startswith('Transaction already closed'): 158 ↛ 160line 158 didn't jump to line 160, because the condition on line 158 was never false

159 raise prisma_errors.TransactionExpiredError(base_error_message) 

160 raise prisma_errors.TransactionError(message) 

161 

162 if 'A value is required but not set' in message: 

163 raise prisma_errors.MissingRequiredValueError(error) 

164 

165 exc: type[Exception] | None = None 

166 

167 kind = user_facing.get('meta', {}).get('kind') 

168 if kind is not None: 

169 exc = META_ERROR_MAPPING.get(kind) 

170 

171 if exc is None: 

172 exc = ERROR_MAPPING.get(code) 

173 

174 if exc is not None: 

175 raise exc(error) 

176 except (KeyError, TypeError) as err: 

177 log.debug('Ignoring error while constructing specialized error %s', err) 

178 continue 

179 

180 try: 

181 raise prisma_errors.DataError(data[0]) 

182 except (IndexError, TypeError): 

183 pass 

184 

185 raise errors.EngineRequestError(resp, f'Could not process erroneous response: {data}')