localstack-python-client/localstack_client/patch.py at refs/heads/master · localstack/localstack-python-client · GitHub

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