Skip to content

Commit c63962f

Browse files
kesmit13claude
andauthored
Add additional check for the workspace endpoint being available (#102)
* feature: Loop on active workspace until endpoint is ready When creating a new workspace with wait_on_active=True, the system now waits for the workspace state to become Active, and then verifies that the endpoint is actually ready by attempting to establish a connection. This ensures that the workspace is truly ready for use before returning to the caller. Added _wait_on_endpoint method in WorkspaceManager that polls the workspace endpoint with configurable interval and timeout, using a context manager to properly handle database connections. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Add debugging messages * Add debugging messages * Add debugging messages * Add debugging messages * Clean up --------- Co-authored-by: Claude <noreply@anthropic.com>
1 parent 5f9cada commit c63962f

File tree

2 files changed

+61
-0
lines changed

2 files changed

+61
-0
lines changed

singlestoredb/management/manager.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
from .. import config
1717
from ..exceptions import ManagementError
18+
from ..exceptions import OperationalError
1819
from .utils import get_token
1920

2021

@@ -310,3 +311,57 @@ def _wait_on_state(
310311
out = getattr(self, f'get_{self.obj_type}')(out.id)
311312

312313
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

singlestoredb/management/workspace.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1794,6 +1794,12 @@ def create_workspace(
17941794
interval=wait_interval,
17951795
timeout=wait_timeout,
17961796
)
1797+
# After workspace is active, wait for endpoint to be ready
1798+
out = self._wait_on_endpoint(
1799+
out,
1800+
interval=wait_interval,
1801+
timeout=wait_timeout,
1802+
)
17971803
return out
17981804

17991805
def get_workspace_group(self, id: str) -> WorkspaceGroup:

0 commit comments

Comments
 (0)