Repository: surrealdb/surrealdb.go
Pull request #256 (2025-07-21 10:16:59)

Enhance Connect, Close, Use, Let, and Unset to accept context.Context
This is a breaking change that enhances the following functions for the `Connection` interface to accept `context.Context`:

- `Connect`
- `Close`
- `Use`
- `Let`
- `Unset`

This also introduces the new top-level function `func Connect(ctx context.Context, connectionURL string) (*DB, error)` to replace the existing, and now deprecated, `func New(connectionURL string) (*DB, error)`.

The existing `New` had a chance of blocking when the initial connection attempt was delayed due to, e.g., unreliable network, and having a context enables the consumer of this SDK to have more control over what to do when it took more time than the designed timeout. Additionally, the new `Connect` function should better convey that it is not a simple local-only operation that initializes the in-memory data of the client/connection, but rather incurs a network call before returning.

The behaviors of `HTTPConnection`'s `Close`, `Use`, `Let`, and `Unset` are unchanged, because they don't incur network calls to cancel. `Connect` is enhanced to respect the provided context when doing an initial health check against SurrealDB.

The behaviors of `WebSocketConnection` counterparts mostly respect the provided context now, allowing you to prevent unexpected blocking when the underlying WebSocket connection is unreliable. The only exception is WebSocketConnection `Close`, which currently ignores `context.Context`. I'll revisit it and probably add a proper context support to it later.

Follow-up to #253 and #254
Ref #100


Patch:
diff --git a/README.md b/README.md
index ed2e7d5..f9b9af5 100644
--- a/README.md
+++ b/README.md
@@ -68,7 +68,7 @@ type Person struct {
 
 func main() {
 	// Connect to SurrealDB
-	db, err := surrealdb.New("ws://localhost:8000")
+	db, err := surrealdb.Connect(context.Background(), "ws://localhost:8000")
 	if err != nil {
 		panic(err)
 	}
diff --git a/db.go b/db.go
index cdbe417..3673436 100644
--- a/db.go
+++ b/db.go
@@ -29,7 +29,22 @@ type DB struct {
 }
 
 // New creates a new SurrealDB client.
+//
+// Deprecated: New is deprecated. Use Connect instead to make your
+// application more robust against network issues.
 func New(connectionURL string) (*DB, error) {
+	return Connect(context.Background(), connectionURL)
+}
+
+// Connect creates a new SurrealDB client and connects to the database.
+//
+// This function incurs a network call (currently HTTP request) to the SurrealDB server to check the health of the connection in
+// case of HTTP, or to establish a WebSocket connection in case of WebSocket.
+//
+// The provided `ctx` is used to cancel the connection attempt if needed,
+// so that you control how long you want to block in case the network is not reliable
+// or any other issues like OS network stack issues/settings/etc.
+func Connect(ctx context.Context, connectionURL string) (*DB, error) {
 	u, err := url.ParseRequestURI(connectionURL)
 	if err != nil {
 		return nil, err
@@ -56,7 +71,7 @@ func New(connectionURL string) (*DB, error) {
 		return nil, fmt.Errorf("invalid connection url")
 	}
 
-	err = con.Connect()
+	err = con.Connect(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -75,13 +90,13 @@ func (db *DB) WithContext(ctx context.Context) *DB {
 }
 
 // Close closes the underlying WebSocket connection.
-func (db *DB) Close() error {
-	return db.con.Close()
+func (db *DB) Close(ctx context.Context) error {
+	return db.con.Close(ctx)
 }
 
 // Use is a method to select the namespace and table to use.
 func (db *DB) Use(ctx context.Context, ns, database string) error {
-	return db.con.Use(ns, database)
+	return db.con.Use(ctx, ns, database)
 }
 
 func (db *DB) Info(ctx context.Context) (map[string]interface{}, error) {
@@ -121,7 +136,7 @@ func (db *DB) SignUp(ctx context.Context, authData interface{}) (string, error)
 		return "", err
 	}
 
-	if err := db.con.Let(constants.AuthTokenKey, token.Result); err != nil {
+	if err := db.con.Let(ctx, constants.AuthTokenKey, token.Result); err != nil {
 		return "", err
 	}
 
@@ -159,7 +174,7 @@ func (db *DB) SignIn(ctx context.Context, authData interface{}) (string, error)
 		return "", err
 	}
 
-	if err := db.con.Let(constants.AuthTokenKey, *token.Result); err != nil {
+	if err := db.con.Let(ctx, constants.AuthTokenKey, *token.Result); err != nil {
 		return "", err
 	}
 
@@ -171,7 +186,7 @@ func (db *DB) Invalidate(ctx context.Context) error {
 		return err
 	}
 
-	if err := db.con.Unset(constants.AuthTokenKey); err != nil {
+	if err := db.con.Unset(ctx, constants.AuthTokenKey); err != nil {
 		return err
 	}
 
@@ -183,7 +198,7 @@ func (db *DB) Authenticate(ctx context.Context, token string) error {
 		return err
 	}
 
-	if err := db.con.Let(constants.AuthTokenKey, token); err != nil {
+	if err := db.con.Let(ctx, constants.AuthTokenKey, token); err != nil {
 		return err
 	}
 
@@ -191,11 +206,11 @@ func (db *DB) Authenticate(ctx context.Context, token string) error {
 }
 
 func (db *DB) Let(ctx context.Context, key string, val interface{}) error {
-	return db.con.Let(key, val)
+	return db.con.Let(ctx, key, val)
 }
 
 func (db *DB) Unset(ctx context.Context, key string) error {
-	return db.con.Unset(key)
+	return db.con.Unset(ctx, key)
 }
 
 func (db *DB) Version(ctx context.Context) (*VersionData, error) {
diff --git a/example/connection.go b/example/connection.go
index 108ee35..89e8b58 100644
--- a/example/connection.go
+++ b/example/connection.go
@@ -29,7 +29,7 @@ func getSurrealDBHTTPURL() string {
 }
 
 func newSurrealDBWSConnection(database string, tables ...string) *surrealdb.DB {
-	db, err := surrealdb.New(getSurrealDBWSURL())
+	db, err := surrealdb.Connect(context.Background(), getSurrealDBWSURL())
 	if err != nil {
 		panic(err)
 	}
@@ -38,7 +38,7 @@ func newSurrealDBWSConnection(database string, tables ...string) *surrealdb.DB {
 }
 
 func newSurrealDBHTTPConnection(database string, tables ...string) *surrealdb.DB {
-	db, err := surrealdb.New(getSurrealDBHTTPURL())
+	db, err := surrealdb.Connect(context.Background(), getSurrealDBHTTPURL())
 	if err != nil {
 		panic(err)
 	}
diff --git a/example/main.go b/example/main.go
index cc0e856..90d0766 100644
--- a/example/main.go
+++ b/example/main.go
@@ -12,7 +12,7 @@ import (
 //nolint:funlen
 func main() {
 	// Connect to SurrealDB
-	db, err := surrealdb.New("ws://localhost:8000")
+	db, err := surrealdb.Connect(context.Background(), "ws://localhost:8000")
 	if err != nil {
 		panic(err)
 	}
diff --git a/pkg/connection/connection.go b/pkg/connection/connection.go
index 41e852d..f3d1ad8 100644
--- a/pkg/connection/connection.go
+++ b/pkg/connection/connection.go
@@ -18,8 +18,8 @@ type LiveHandler interface {
 }
 
 type Connection interface {
-	Connect() error
-	Close() error
+	Connect(ctx context.Context) error
+	Close(ctx context.Context) error
 	// Send sends a request to SurrealDB and expects a response.
 	//
 	// It requires `res` to be of type `*RPCResponse[T]` where T is a type that implements `cbor.Unmarshaller`,
@@ -28,9 +28,9 @@ type Connection interface {
 	//
 	// The `ctx` is used to cancel the request if the context is canceled.
 	Send(ctx context.Context, res interface{}, method string, params ...interface{}) error
-	Use(namespace string, database string) error
-	Let(key string, value interface{}) error
-	Unset(key string) error
+	Use(ctx context.Context, namespace string, database string) error
+	Let(ctx context.Context, key string, value interface{}) error
+	Unset(ctx context.Context, key string) error
 	LiveNotifications(id string) (chan Notification, error)
 	GetUnmarshaler() codec.Unmarshaler
 }
diff --git a/pkg/connection/embedded.go b/pkg/connection/embedded.go
index 1b46dc3..d172d23 100644
--- a/pkg/connection/embedded.go
+++ b/pkg/connection/embedded.go
@@ -56,7 +56,7 @@ func NewEmbeddedConnection(p NewConnectionParams) *EmbeddedConnection {
 	return &con
 }
 
-func (h *EmbeddedConnection) Connect() error {
+func (h *EmbeddedConnection) Connect(ctx context.Context) error {
 	err := h.preConnectionChecks()
 	if err != nil {
 		return err
@@ -91,7 +91,7 @@ func (h *EmbeddedConnection) Connect() error {
 	return nil
 }
 
-func (h *EmbeddedConnection) Close() error {
+func (h *EmbeddedConnection) Close(ctx context.Context) error {
 	C.sr_surreal_rpc_free(h.surrealRPC)
 
 	h.surrealRPC = nil
@@ -139,14 +139,14 @@ func (h *EmbeddedConnection) Send(ctx context.Context, res interface{}, method s
 	return h.unmarshaler.Unmarshal(rpcRes, res)
 }
 
-func (h *EmbeddedConnection) Use(namespace, database string) error {
-	return h.Send(context.Background(), nil, "use", namespace, database)
+func (h *EmbeddedConnection) Use(ctx context.Context, namespace, database string) error {
+	return h.Send(ctx, nil, "use", namespace, database)
 }
 
-func (h *EmbeddedConnection) Let(key string, value interface{}) error {
-	return h.Send(context.Background(), nil, "let", key, value)
+func (h *EmbeddedConnection) Let(ctx context.Context, key string, value interface{}) error {
+	return h.Send(ctx, nil, "let", key, value)
 }
 
-func (h *EmbeddedConnection) Unset(key string) error {
-	return h.Send(context.Background(), nil, "unset", key)
+func (h *EmbeddedConnection) Unset(ctx context.Context, key string) error {
+	return h.Send(ctx, nil, "unset", key)
 }
diff --git a/pkg/connection/http.go b/pkg/connection/http.go
index f200fb3..b30ea02 100644
--- a/pkg/connection/http.go
+++ b/pkg/connection/http.go
@@ -41,8 +41,7 @@ func NewHTTPConnection(p NewConnectionParams) *HTTPConnection {
 	return &con
 }
 
-func (h *HTTPConnection) Connect() error {
-	ctx := context.TODO()
+func (h *HTTPConnection) Connect(ctx context.Context) error {
 	if err := h.preConnectionChecks(); err != nil {
 		return err
 	}
@@ -59,7 +58,7 @@ func (h *HTTPConnection) Connect() error {
 	return nil
 }
 
-func (h *HTTPConnection) Close() error {
+func (h *HTTPConnection) Close(ctx context.Context) error {
 	return nil
 }
 
@@ -164,19 +163,19 @@ func (h *HTTPConnection) MakeRequest(req *http.Request) ([]byte, error) {
 	return nil, errorResponse.Error
 }
 
-func (h *HTTPConnection) Use(namespace, database string) error {
+func (h *HTTPConnection) Use(ctx context.Context, namespace, database string) error {
 	h.variables.Store("namespace", namespace)
 	h.variables.Store("database", database)
 
 	return nil
 }
 
-func (h *HTTPConnection) Let(key string, value interface{}) error {
+func (h *HTTPConnection) Let(ctx context.Context, key string, value interface{}) error {
 	h.variables.Store(key, value)
 	return nil
 }
 
-func (h *HTTPConnection) Unset(key string) error {
+func (h *HTTPConnection) Unset(ctx context.Context, key string) error {
 	h.variables.Delete(key)
 	return nil
 }
diff --git a/pkg/connection/ws.go b/pkg/connection/ws.go
index 053da38..f3dc5a0 100644
--- a/pkg/connection/ws.go
+++ b/pkg/connection/ws.go
@@ -76,12 +76,12 @@ func NewWebSocketConnection(p NewConnectionParams) *WebSocketConnection {
 	}
 }
 
-func (ws *WebSocketConnection) Connect() error {
+func (ws *WebSocketConnection) Connect(ctx context.Context) error {
 	if err := ws.preConnectionChecks(); err != nil {
 		return err
 	}
 
-	connection, res, err := DefaultDialer.Dial(fmt.Sprintf("%s/rpc", ws.baseURL), nil)
+	connection, res, err := DefaultDialer.DialContext(ctx, fmt.Sprintf("%s/rpc", ws.baseURL), nil)
 	if err != nil {
 		return err
 	}
@@ -126,7 +126,11 @@ func (ws *WebSocketConnection) SetCompression(compress bool) *WebSocketConnectio
 	return ws
 }
 
-func (ws *WebSocketConnection) Close() error {
+// Close closes the WebSocket connection and stops listening for incoming messages.
+//
+// Note that this method may block until the Close message is sent,
+// even though the provided context is canceled beforehand, for now.
+func (ws *WebSocketConnection) Close(ctx context.Context) error {
 	ws.connLock.Lock()
 	defer ws.connLock.Unlock()
 	close(ws.closeChan)
@@ -138,16 +142,16 @@ func (ws *WebSocketConnection) Close() error {
 	return ws.Conn.Close()
 }
 
-func (ws *WebSocketConnection) Use(namespace, database string) error {
-	return ws.Send(context.Background(), nil, "use", namespace, database)
+func (ws *WebSocketConnection) Use(ctx context.Context, namespace, database string) error {
+	return ws.Send(ctx, nil, "use", namespace, database)
 }
 
-func (ws *WebSocketConnection) Let(key string, value interface{}) error {
-	return ws.Send(context.Background(), nil, "let", key, value)
+func (ws *WebSocketConnection) Let(ctx context.Context, key string, value interface{}) error {
+	return ws.Send(ctx, nil, "let", key, value)
 }
 
-func (ws *WebSocketConnection) Unset(key string) error {
-	return ws.Send(context.Background(), nil, "unset", key)
+func (ws *WebSocketConnection) Unset(ctx context.Context, key string) error {
+	return ws.Send(ctx, nil, "unset", key)
 }
 
 func (ws *WebSocketConnection) GetUnmarshaler() codec.Unmarshaler {
@@ -177,9 +181,6 @@ func (ws *WebSocketConnection) Send(ctx context.Context, dest interface{}, metho
 	case <-ws.closeChan:
 		return ws.closeError
 	case <-ctx.Done():
-		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-			return constants.ErrTimeout
-		}
 		return ctx.Err()
 	default:
 	}
@@ -203,9 +204,6 @@ func (ws *WebSocketConnection) Send(ctx context.Context, dest interface{}, metho
 
 	select {
 	case <-ctx.Done():
-		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-			return constants.ErrTimeout
-		}
 		return ctx.Err()
 	case res, open := <-responseChan:
 		if !open {
diff --git a/pkg/constants/errors.go b/pkg/constants/errors.go
index 8c67815..2a3d963 100644
--- a/pkg/constants/errors.go
+++ b/pkg/constants/errors.go
@@ -10,7 +10,6 @@ var (
 )
 var (
 	ErrIDInUse            = errors.New("id already in use")
-	ErrTimeout            = errors.New("timeout")
 	ErrNoBaseURL          = errors.New("base url not set")
 	ErrNoMarshaler        = errors.New("marshaler is not set")
 	ErrNoUnmarshaler      = errors.New("unmarshaler is not set")


Test patch:
diff --git a/db_test.go b/db_test.go
index bb0d2f0..e5583db 100644
--- a/db_test.go
+++ b/db_test.go
@@ -84,7 +84,7 @@ func (s *SurrealDBTestSuite) TearDownTest() {
 
 // TearDownSuite is called after the s has finished running
 func (s *SurrealDBTestSuite) TearDownSuite() {
-	err := s.db.Close()
+	err := s.db.Close(context.Background())
 	s.Require().NoError(err)
 }
 
diff --git a/internal/benchmark/benchmark_test.go b/internal/benchmark/benchmark_test.go
index 8c45cf5..eab3505 100644
--- a/internal/benchmark/benchmark_test.go
+++ b/internal/benchmark/benchmark_test.go
@@ -18,7 +18,7 @@ type testUser struct {
 }
 
 func SetupMockDB() (*surrealdb.DB, error) {
-	return surrealdb.New("")
+	return surrealdb.Connect(context.Background(), "")
 }
 
 func BenchmarkCreate(b *testing.B) {
diff --git a/pkg/connection/connection_test.go b/pkg/connection/connection_test.go
index 286f313..d4d0d18 100644
--- a/pkg/connection/connection_test.go
+++ b/pkg/connection/connection_test.go
@@ -57,11 +57,11 @@ func (s *ConnectionTestSuite) SetupSuite() {
 	con := s.connImplementations[s.name]
 
 	// connect
-	err := con.Connect()
+	err := con.Connect(context.Background())
 	s.Require().NoError(err)
 
 	// set namespace, database
-	err = con.Use("test", "test")
+	err = con.Use(context.Background(), "test", "test")
 	s.Require().NoError(err)
 
 	// sign in
@@ -71,12 +71,12 @@ func (s *ConnectionTestSuite) SetupSuite() {
 		"pass": "root",
 	})
 	s.Require().NoError(err)
-	_ = con.Let(constants.AuthTokenKey, *token.Result)
+	_ = con.Let(context.Background(), constants.AuthTokenKey, *token.Result)
 }
 
 func (s *ConnectionTestSuite) TearDownSuite() {
 	con := s.connImplementations[s.name]
-	err := con.Close()
+	err := con.Close(context.Background())
 	s.Require().NoError(err)
 }