18
18
19
19
import logging
20
20
import warnings
21
- from typing import Optional
21
+ from typing import Any , Callable , Dict , Optional , Union
22
22
23
- from elastic_transport import AsyncTransport , TransportError
23
+ from elastic_transport import AsyncTransport , NodeConfig , TransportError
24
24
from elastic_transport .client_utils import DEFAULT
25
25
26
26
from ...exceptions import NotFoundError
27
27
from ...serializer import DEFAULT_SERIALIZERS
28
- from ._base import BaseClient , resolve_auth_headers
28
+ from ._base import (
29
+ BaseClient ,
30
+ create_sniff_callback ,
31
+ default_sniff_callback ,
32
+ resolve_auth_headers ,
33
+ )
29
34
from .async_search import AsyncSearchClient
30
35
from .autoscaling import AutoscalingClient
31
36
from .cat import CatClient
@@ -148,9 +153,21 @@ def __init__(
148
153
sniff_on_node_failure = DEFAULT ,
149
154
sniff_timeout = DEFAULT ,
150
155
min_delay_between_sniffing = DEFAULT ,
156
+ sniffed_node_callback : Optional [
157
+ Callable [[Dict [str , Any ], NodeConfig ], Optional [NodeConfig ]]
158
+ ] = None ,
151
159
meta_header = DEFAULT ,
152
160
# Deprecated
153
161
timeout = DEFAULT ,
162
+ randomize_hosts = DEFAULT ,
163
+ host_info_callback : Optional [
164
+ Callable [
165
+ [Dict [str , Any ], Dict [str , Union [str , int ]]],
166
+ Optional [Dict [str , Union [str , int ]]],
167
+ ]
168
+ ] = None ,
169
+ sniffer_timeout = DEFAULT ,
170
+ sniff_on_connection_fail = DEFAULT ,
154
171
# Internal use only
155
172
_transport : Optional [AsyncTransport ] = None ,
156
173
) -> None :
@@ -170,6 +187,86 @@ def __init__(
170
187
)
171
188
request_timeout = timeout
172
189
190
+ if randomize_hosts is not DEFAULT :
191
+ if randomize_nodes_in_pool is not DEFAULT :
192
+ raise ValueError (
193
+ "Can't specify both 'randomize_hosts' and 'randomize_nodes_in_pool', "
194
+ "instead only specify 'randomize_nodes_in_pool'"
195
+ )
196
+ warnings .warn (
197
+ "The 'randomize_hosts' parameter is deprecated in favor of 'randomize_nodes_in_pool'" ,
198
+ category = DeprecationWarning ,
199
+ stacklevel = 2 ,
200
+ )
201
+ randomize_nodes_in_pool = randomize_hosts
202
+
203
+ if sniffer_timeout is not DEFAULT :
204
+ if min_delay_between_sniffing is not DEFAULT :
205
+ raise ValueError (
206
+ "Can't specify both 'sniffer_timeout' and 'min_delay_between_sniffing', "
207
+ "instead only specify 'min_delay_between_sniffing'"
208
+ )
209
+ warnings .warn (
210
+ "The 'sniffer_timeout' parameter is deprecated in favor of 'min_delay_between_sniffing'" ,
211
+ category = DeprecationWarning ,
212
+ stacklevel = 2 ,
213
+ )
214
+ min_delay_between_sniffing = sniffer_timeout
215
+
216
+ if sniff_on_connection_fail is not DEFAULT :
217
+ if sniff_on_node_failure is not DEFAULT :
218
+ raise ValueError (
219
+ "Can't specify both 'sniff_on_connection_fail' and 'sniff_on_node_failure', "
220
+ "instead only specify 'sniff_on_node_failure'"
221
+ )
222
+ warnings .warn (
223
+ "The 'sniff_on_connection_fail' parameter is deprecated in favor of 'sniff_on_node_failure'" ,
224
+ category = DeprecationWarning ,
225
+ stacklevel = 2 ,
226
+ )
227
+ sniff_on_node_failure = sniff_on_connection_fail
228
+
229
+ # Setting min_delay_between_sniffing=True implies sniff_before_requests=True
230
+ if min_delay_between_sniffing is not DEFAULT :
231
+ sniff_before_requests = True
232
+
233
+ sniffing_options = (
234
+ sniff_timeout ,
235
+ sniff_on_start ,
236
+ sniff_before_requests ,
237
+ sniff_on_node_failure ,
238
+ sniffed_node_callback ,
239
+ min_delay_between_sniffing ,
240
+ sniffed_node_callback ,
241
+ )
242
+ if cloud_id is not None and any (
243
+ x is not DEFAULT and x is not None for x in sniffing_options
244
+ ):
245
+ raise ValueError (
246
+ "Sniffing should not be enabled when connecting to Elastic Cloud"
247
+ )
248
+
249
+ sniff_callback = None
250
+ if host_info_callback is not None :
251
+ if sniffed_node_callback is not None :
252
+ raise ValueError (
253
+ "Can't specify both 'host_info_callback' and 'sniffed_node_callback', "
254
+ "instead only specify 'sniffed_node_callback'"
255
+ )
256
+ sniff_callback = create_sniff_callback (
257
+ host_info_callback = host_info_callback
258
+ )
259
+ elif sniffed_node_callback is not None :
260
+ sniff_callback = create_sniff_callback (
261
+ sniffed_node_callback = sniffed_node_callback
262
+ )
263
+ elif (
264
+ sniff_on_start is True
265
+ or sniff_before_requests is True
266
+ or sniff_on_node_failure is True
267
+ ):
268
+ sniff_callback = default_sniff_callback
269
+
173
270
if _transport is None :
174
271
node_configs = client_node_configs (
175
272
hosts ,
@@ -222,6 +319,7 @@ def __init__(
222
319
_transport = transport_class (
223
320
node_configs ,
224
321
client_meta_service = CLIENT_META_SERVICE ,
322
+ sniff_callback = sniff_callback ,
225
323
** transport_kwargs ,
226
324
)
227
325
0 commit comments