-
Notifications
You must be signed in to change notification settings - Fork 30
Expand file tree
/
Copy pathpatch.py
More file actions
151 lines (131 loc) · 4.53 KB
/
patch.py
File metadata and controls
151 lines (131 loc) · 4.53 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import types
import boto3
from boto3.session import Session
from botocore.serialize import Serializer
_state = {}
DEFAULT_ACCESS_KEY_ID = "test"
DEFAULT_SECRET_ACCESS_KEY = "test"
def enable_local_endpoints():
"""Patch the boto3 library to transparently use the LocalStack endpoints by default."""
from localstack_client.config import get_service_endpoint
def _add_custom_kwargs(
kwargs,
service_name,
endpoint_url=None,
aws_access_key_id=None,
aws_secret_access_key=None,
):
kwargs["endpoint_url"] = endpoint_url or get_service_endpoint(service_name)
kwargs["aws_access_key_id"] = aws_access_key_id or DEFAULT_ACCESS_KEY_ID
kwargs["aws_secret_access_key"] = (
aws_secret_access_key or DEFAULT_SECRET_ACCESS_KEY
)
def _client(
self,
service_name,
region_name=None,
api_version=None,
use_ssl=True,
verify=None,
endpoint_url=None,
aws_access_key_id=None,
aws_secret_access_key=None,
**kwargs,
):
_add_custom_kwargs(
kwargs,
service_name,
endpoint_url=endpoint_url,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
)
return _client_orig(
self,
service_name,
region_name=region_name,
api_version=api_version,
use_ssl=use_ssl,
verify=verify,
**kwargs,
)
def _resource(
self,
service_name,
region_name=None,
api_version=None,
use_ssl=True,
verify=None,
endpoint_url=None,
aws_access_key_id=None,
aws_secret_access_key=None,
**kwargs,
):
_add_custom_kwargs(
kwargs,
service_name,
endpoint_url=endpoint_url,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
)
return _resource_orig(
self,
service_name,
region_name=region_name,
api_version=api_version,
use_ssl=use_ssl,
verify=verify,
**kwargs,
)
if _state.get("_client_orig"):
# patch already applied -> return
return
# patch boto3 default session (if available)
try:
session = boto3._get_default_session()
_state["_default_client_orig"] = session.client
session.client = types.MethodType(_client, session)
_state["_default_resource_orig"] = session.resource
session.resource = types.MethodType(_resource, session)
except Exception:
# swallowing for now - looks like the default session is not available (yet)
pass
# patch session.client(..)
_client_orig = Session.client
_state["_client_orig"] = _client_orig
Session.client = _client
# patch session.resource(..)
_resource_orig = Session.resource
_state["_resource_orig"] = _resource_orig
Session.resource = _resource
def disable_local_endpoints():
"""Disable the boto3 patches and revert to using the default endpoints against real AWS."""
_client = _state.pop("_client_orig", None)
if _client:
Session.client = _client
_resource = _state.pop("_resource_orig", None)
if _resource:
Session.resource = _resource
# undo patches for boto3 default session
try:
session = boto3._get_default_session()
if _state.get("_default_client_orig"):
session.client = _state["_default_client_orig"]
if _state.get("_default_resource_orig"):
session.resource = _state["_default_resource_orig"]
except Exception:
pass
def patch_expand_host_prefix():
"""Apply a patch to botocore, to skip adding host prefixes to endpoint URLs"""
def _expand_host_prefix(self, parameters, operation_model, *args, **kwargs):
result = _expand_host_prefix_orig(
self, parameters, operation_model, *args, **kwargs
)
# skip adding host prefixes, to avoid making requests to, e.g., http://data-localhost:4566
is_sd = operation_model.service_model.service_name == "servicediscovery"
if is_sd and result == "data-":
return None
if operation_model.service_model.service_name == "mwaa" and result == "api.":
return None
return result
_expand_host_prefix_orig = Serializer._expand_host_prefix
Serializer._expand_host_prefix = _expand_host_prefix