diff --git a/singlestoredb/management/manager.py b/singlestoredb/management/manager.py index 3d974950..076d0716 100644 --- a/singlestoredb/management/manager.py +++ b/singlestoredb/management/manager.py @@ -15,6 +15,7 @@ from .. import config from ..exceptions import ManagementError +from ..exceptions import OperationalError from .utils import get_token @@ -310,3 +311,57 @@ def _wait_on_state( out = getattr(self, f'get_{self.obj_type}')(out.id) return out + + def _wait_on_endpoint( + self, + out: Any, + interval: int = 10, + timeout: int = 300, + ) -> Any: + """ + Wait for the endpoint to be ready by attempting to connect. + + Parameters + ---------- + out : Any + Workspace object with a connect method + interval : int, optional + Interval between each connection attempt (default: 10 seconds) + timeout : int, optional + Maximum time to wait before raising an exception (default: 300 seconds) + + Raises + ------ + ManagementError + If timeout is reached or endpoint is not available + + Returns + ------- + Same object type as `out` + + """ + if not hasattr(out, 'connect') or not out.connect: + raise ManagementError( + msg=f'{type(out).__name__} object does not have a valid endpoint', + ) + + while True: + try: + # Try to establish a connection to the endpoint using context manager + with out.connect(connect_timeout=5): + pass + except Exception as exc: + # If we get an 'access denied' error, that means that the server is + # up and we just aren't authenticating. + if isinstance(exc, OperationalError) and exc.errno == 1045: + break + # If connection fails, check timeout and retry + if timeout <= 0: + raise ManagementError( + msg=f'Exceeded waiting time for {self.obj_type} endpoint ' + 'to become ready', + ) + time.sleep(interval) + timeout -= interval + + return out diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py index 7c394ade..b92dc833 100644 --- a/singlestoredb/management/workspace.py +++ b/singlestoredb/management/workspace.py @@ -1794,6 +1794,12 @@ def create_workspace( interval=wait_interval, timeout=wait_timeout, ) + # After workspace is active, wait for endpoint to be ready + out = self._wait_on_endpoint( + out, + interval=wait_interval, + timeout=wait_timeout, + ) return out def get_workspace_group(self, id: str) -> WorkspaceGroup: