|
15 | 15 |
|
16 | 16 | from .. import config |
17 | 17 | from ..exceptions import ManagementError |
| 18 | +from ..exceptions import OperationalError |
18 | 19 | from .utils import get_token |
19 | 20 |
|
20 | 21 |
|
@@ -310,3 +311,57 @@ def _wait_on_state( |
310 | 311 | out = getattr(self, f'get_{self.obj_type}')(out.id) |
311 | 312 |
|
312 | 313 | return out |
| 314 | + |
| 315 | + def _wait_on_endpoint( |
| 316 | + self, |
| 317 | + out: Any, |
| 318 | + interval: int = 10, |
| 319 | + timeout: int = 300, |
| 320 | + ) -> Any: |
| 321 | + """ |
| 322 | + Wait for the endpoint to be ready by attempting to connect. |
| 323 | +
|
| 324 | + Parameters |
| 325 | + ---------- |
| 326 | + out : Any |
| 327 | + Workspace object with a connect method |
| 328 | + interval : int, optional |
| 329 | + Interval between each connection attempt (default: 10 seconds) |
| 330 | + timeout : int, optional |
| 331 | + Maximum time to wait before raising an exception (default: 300 seconds) |
| 332 | +
|
| 333 | + Raises |
| 334 | + ------ |
| 335 | + ManagementError |
| 336 | + If timeout is reached or endpoint is not available |
| 337 | +
|
| 338 | + Returns |
| 339 | + ------- |
| 340 | + Same object type as `out` |
| 341 | +
|
| 342 | + """ |
| 343 | + if not hasattr(out, 'connect') or not out.connect: |
| 344 | + raise ManagementError( |
| 345 | + msg=f'{type(out).__name__} object does not have a valid endpoint', |
| 346 | + ) |
| 347 | + |
| 348 | + while True: |
| 349 | + try: |
| 350 | + # Try to establish a connection to the endpoint using context manager |
| 351 | + with out.connect(connect_timeout=5): |
| 352 | + pass |
| 353 | + except Exception as exc: |
| 354 | + # If we get an 'access denied' error, that means that the server is |
| 355 | + # up and we just aren't authenticating. |
| 356 | + if isinstance(exc, OperationalError) and exc.errno == 1045: |
| 357 | + break |
| 358 | + # If connection fails, check timeout and retry |
| 359 | + if timeout <= 0: |
| 360 | + raise ManagementError( |
| 361 | + msg=f'Exceeded waiting time for {self.obj_type} endpoint ' |
| 362 | + 'to become ready', |
| 363 | + ) |
| 364 | + time.sleep(interval) |
| 365 | + timeout -= interval |
| 366 | + |
| 367 | + return out |
0 commit comments