Add additional check for the workspace endpoint being available

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

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: