Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/build/lib/airflow/api_internal/internal_api_call.py: 43%

1# Licensed to the Apache Software Foundation (ASF) under one 

2# or more contributor license agreements. See the NOTICE file 

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

5# to you under the Apache License, Version 2.0 (the 

6# "License"); you may not use this file except in compliance 

7# with the License. You may obtain a copy of the License at 

8# 

9# http://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, 

12# software distributed under the License is distributed on an 

13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

14# KIND, either express or implied. See the License for the 

15# specific language governing permissions and limitations 

16# under the License. 

17 

18from __future__ import annotations 

19 

20import inspect 

21import json 

22from functools import wraps 

23from typing import Callable, TypeVar 

24 

25import requests 

26 

27from airflow.configuration import conf 

28from airflow.exceptions import AirflowConfigException, AirflowException 

29from airflow.settings import _ENABLE_AIP_44 

30from airflow.typing_compat import ParamSpec 

31 

32PS = ParamSpec("PS") 

33RT = TypeVar("RT") 

34 

35 

36class InternalApiConfig: 

37 """Stores and caches configuration for Internal API.""" 

38 

39 _initialized = False 

40 _use_internal_api = False 

41 _internal_api_endpoint = "" 

42 

43 @staticmethod 

44 def force_database_direct_access(): 

45 """Current component will not use Internal API. 

46 

47 All methods decorated with internal_api_call will always be executed locally. 

48 This mode is needed for "trusted" components like Scheduler, Webserver or Internal Api server. 

49 """ 

50 InternalApiConfig._initialized = True 

51 InternalApiConfig._use_internal_api = False 

52 

53 @staticmethod 

54 def get_use_internal_api(): 

55 if not InternalApiConfig._initialized: 

56 InternalApiConfig._init_values() 

57 return InternalApiConfig._use_internal_api 

58 

59 @staticmethod 

60 def get_internal_api_endpoint(): 

61 if not InternalApiConfig._initialized: 

62 InternalApiConfig._init_values() 

63 return InternalApiConfig._internal_api_endpoint 

64 

65 @staticmethod 

66 def _init_values(): 

67 use_internal_api = conf.getboolean("core", "database_access_isolation", fallback=False) 

68 if use_internal_api and not _ENABLE_AIP_44: 

69 raise RuntimeError("The AIP_44 is not enabled so you cannot use it.") 

70 internal_api_endpoint = "" 

71 if use_internal_api: 

72 internal_api_url = conf.get("core", "internal_api_url") 

73 internal_api_endpoint = internal_api_url + "/internal_api/v1/rpcapi" 

74 if not internal_api_endpoint.startswith("http://"): 

75 raise AirflowConfigException("[core]internal_api_url must start with http://") 

76 

77 InternalApiConfig._initialized = True 

78 InternalApiConfig._use_internal_api = use_internal_api 

79 InternalApiConfig._internal_api_endpoint = internal_api_endpoint 

80 

81 

82def internal_api_call(func: Callable[PS, RT]) -> Callable[PS, RT]: 

83 """Decorator for methods which may be executed in database isolation mode. 

84 

85 If [core]database_access_isolation is true then such method are not executed locally, 

86 but instead RPC call is made to Database API (aka Internal API). This makes some components 

87 decouple from direct Airflow database access. 

88 Each decorated method must be present in METHODS list in airflow.api_internal.endpoints.rpc_api_endpoint. 

89 Only static methods can be decorated. This decorator must be before "provide_session". 

90 

91 See [AIP-44](https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-44+Airflow+Internal+API) 

92 for more information . 

93 """ 

94 headers = { 

95 "Content-Type": "application/json", 

96 } 

97 

98 def make_jsonrpc_request(method_name: str, params_json: str) -> bytes: 

99 data = {"jsonrpc": "2.0", "method": method_name, "params": params_json} 

100 internal_api_endpoint = InternalApiConfig.get_internal_api_endpoint() 

101 response = requests.post(url=internal_api_endpoint, data=json.dumps(data), headers=headers) 

102 if response.status_code != 200: 

103 raise AirflowException( 

104 f"Got {response.status_code}:{response.reason} when sending the internal api request." 

105 ) 

106 return response.content 

107 

108 @wraps(func) 

109 def wrapper(*args, **kwargs) -> RT: 

110 use_internal_api = InternalApiConfig.get_use_internal_api() 

111 if not use_internal_api: 

112 return func(*args, **kwargs) 

113 

114 from airflow.serialization.serialized_objects import BaseSerialization # avoid circular import 

115 

116 bound = inspect.signature(func).bind(*args, **kwargs) 

117 arguments_dict = dict(bound.arguments) 

118 if "session" in arguments_dict: 

119 del arguments_dict["session"] 

120 if "cls" in arguments_dict: # used by @classmethod 

121 del arguments_dict["cls"] 

122 

123 args_json = json.dumps(BaseSerialization.serialize(arguments_dict, use_pydantic_models=True)) 

124 method_name = f"{func.__module__}.{func.__qualname__}" 

125 result = make_jsonrpc_request(method_name, args_json) 

126 return BaseSerialization.deserialize(json.loads(result), use_pydantic_models=True) 

127 

128 return wrapper