Added NewWithOptions func

db.go

@@ -32,6 +32,16 @@ func New(url string) (*DB, error) {
  return &DB{ws}, nil
 }
 
+// NewWithOptions Creates a new DB instance given a WebSocket URL and options.
+// Options are passed to the websocket library.
+func NewWithOptions(url string, options ...websocket.Option) (*DB, error) {
+ ws, err := websocket.NewWebsocketWithOptions(url, options...)
+ if err != nil {
+ return nil, err
+ }
+ return &DB{ws}, nil
+}
+
 // Unmarshal loads a SurrealDB response into a struct.
 func Unmarshal(data, v interface{}) error {
  var jsonBytes []byte

db_test.go

@@ -4,10 +4,12 @@ import (
  "fmt"
  "os"
  "testing"
+ "time"
 
  "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/require"
  "github.com/surrealdb/surrealdb.go"
+ "github.com/surrealdb/surrealdb.go/internal/websocket"
 )
 
 // a simple user struct for testing
@@ -24,6 +26,18 @@ func (t testUser) String() string {
  return fmt.Sprintf("testUser{Username: %+v, Password: %+v, ID: %+v}", t.Username, t.Password, t.ID)
 }
 
+func openConnection(t *testing.T) *surrealdb.DB {
+ url := os.Getenv("SURREALDB_URL")
+ if url == "" {
+ url = "ws://localhost:8000/rpc"
+ }
+
+ db, err := surrealdb.New(url)
+ require.NoError(t, err)
+
+ return db
+}
+
 func setupDB(t *testing.T) *surrealdb.DB {
  db := openConnection(t)
  _ = signin(t, db)
@@ -32,17 +46,37 @@ func setupDB(t *testing.T) *surrealdb.DB {
  return db
 }
 
-func openConnection(t *testing.T) *surrealdb.DB {
+func setupDBOptions(t *testing.T) *surrealdb.DB {
+ db := openConnectionWithOptions(t)
+ _ = signin(t, db)
+ _, err := db.Use("test", "test")
+ require.NoError(t, err)
+ return db
+}
+
+func openConnectionWithOptions(t *testing.T) *surrealdb.DB {
  url := os.Getenv("SURREALDB_URL")
  if url == "" {
  url = "ws://localhost:8000/rpc"
  }
 
- db, err := surrealdb.New(url)
- if err != nil {
- t.Fatal(err)
+ option1 := func(ws *websocket.WebSocket) error {
+  ws.Timeout = time.Duration(20) * time.Second
+ return nil
  }
 
+ option2 := func(ws *websocket.WebSocket) error {
+ ws.Conn.EnableWriteCompression(true)
+ return nil
+ }
+
+ options := []websocket.Option{option1, option2}
+
+ db, err := surrealdb.NewWithOptions(url, options...)
+
+ require.NotNil(t, db)
+ require.NoError(t, err)
+
  return db
 }
 
@@ -52,22 +86,35 @@ func signin(t *testing.T, db *surrealdb.DB) interface{} {
  "pass": "root",
  })
 
- if err != nil {
- t.Fatal(err)
- }
+ require.NoError(t, err)
 
  return signin
 }
 
 func TestDelete(t *testing.T) {
- db := openConnection(t)
+ db := setupDB(t)
  defer db.Close()
 
- _ = signin(t, db)
+ userData, err := db.Create("users", testUser{
+ Username: "johnny",
+ Password: "123",
+ })
+ require.NoError(t, err)
 
- _, err := db.Use("test", "test")
+ // unmarshal the data into a user struct
+ var user []testUser
+ err = surrealdb.Unmarshal(userData, &user)
  require.NoError(t, err)
 
+ // Delete the users...
+ _, err = db.Delete("users")
+ require.NoError(t, err)
+}
+
+func TestDeleteWithOptions(t *testing.T) {
+ db := setupDBOptions(t)
+ defer db.Close()
+
  userData, err := db.Create("users", testUser{
  Username: "johnny",
  Password: "123",
@@ -85,14 +132,9 @@ func TestDelete(t *testing.T) {
 }
 
 func TestCreate(t *testing.T) {
- db := openConnection(t)
+ db := setupDB(t)
  defer db.Close()
 
- _ = signin(t, db)
-
- _, err := db.Use("test", "test")
- require.NoError(t, err)
-
  t.Run("raw map works", func(t *testing.T) {
  userData, err := db.Create("users", map[string]interface{}{
  "username": "johnny",
@@ -153,14 +195,9 @@ func TestCreate(t *testing.T) {
 }
 
 func TestSelect(t *testing.T) {
- db := openConnection(t)
+ db := setupDB(t)
  defer db.Close()
 
- _ = signin(t, db)
-
- _, err := db.Use("test", "test")
- require.NoError(t, err)
-
  createdUsers, err := db.Create("users", testUser{
  Username: "johnnyjohn",
  Password: "123",
@@ -201,14 +238,9 @@ func TestSelect(t *testing.T) {
 }
 
 func TestUpdate(t *testing.T) {
- db := openConnection(t)
+ db := setupDB(t)
  defer db.Close()
 
- _ = signin(t, db)
-
- _, err := db.Use("test", "test")
- require.NoError(t, err)
-
  userData, err := db.Create("users", testUser{
  Username: "johnny",
  Password: "123",
@@ -237,15 +269,10 @@ func TestUpdate(t *testing.T) {
 }
 
 func TestUnmarshalRaw(t *testing.T) {
- db := openConnection(t)
+ db := setupDB(t)
  defer db.Close()
 
- _ = signin(t, db)
-
- _, err := db.Use("test", "test")
- require.NoError(t, err)
-
- _, err = db.Delete("users")
+ _, err := db.Delete("users")
  require.NoError(t, err)
 
  username := "johnny"
@@ -278,15 +305,10 @@ func TestUnmarshalRaw(t *testing.T) {
 }
 
 func TestModify(t *testing.T) {
- db := openConnection(t)
+ db := setupDB(t)
  defer db.Close()
 
- _ = signin(t, db)
-
- _, err := db.Use("test", "test")
- require.NoError(t, err)
-
- _, err = db.Delete("users:999") // Cleanup for reproducibility
+ _, err := db.Delete("users:999") // Cleanup for reproducibility
  require.NoError(t, err)
 
  _, err = db.Create("users:999", map[string]interface{}{
@@ -313,21 +335,16 @@ func TestModify(t *testing.T) {
 }
 
 func TestNonRowSelect(t *testing.T) {
- db := openConnection(t)
+ db := setupDB(t)
  defer db.Close()
 
- _ = signin(t, db)
-
- _, err := db.Use("test", "test")
- require.NoError(t, err)
-
  user := testUser{
  Username: "ElecTwix",
  Password: "1234",
  ID: "users:notexists",
  }
 
- _, err = db.Select("users:notexists")
+ _, err := db.Select("users:notexists")
  assert.Equal(t, err, surrealdb.ErrNoRow)
 
  _, err = surrealdb.SmartUnmarshal[testUser](db.Select("users:notexists"))

internal/websocket/ws.go

@@ -23,9 +23,9 @@ const (
 type Option func(ws *WebSocket) error
 
 type WebSocket struct {
- conn *websocket.Conn
+ Conn *websocket.Conn
  connLock sync.Mutex
- timeout time.Duration
+ Timeout time.Duration
 
  responseChannels map[string]chan RPCResponse
  responseChannelsLock sync.RWMutex
@@ -47,7 +47,7 @@ func NewWebsocketWithOptions(url string, options ...Option) (*WebSocket, error)
  }
 
  ws := &WebSocket{
- conn: conn,
+ Conn: conn,
  close: make(chan int),
  responseChannels: make(map[string]chan RPCResponse),
  }
@@ -64,7 +64,7 @@ func NewWebsocketWithOptions(url string, options ...Option) (*WebSocket, error)
 
 func Timeout(timeout float64) Option {
  return func(ws *WebSocket) error {
- ws.timeout = time.Duration(timeout) * time.Second
+ ws.Timeout = time.Duration(timeout) * time.Second
  return nil
  }
 }
@@ -74,7 +74,7 @@ func (ws *WebSocket) Close() error {
  close(ws.close)
  }()
 
- return ws.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(CloseMessageCode, ""))
+ return ws.Conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(CloseMessageCode, ""))
 }
 
 var (
@@ -128,7 +128,7 @@ func (ws *WebSocket) Send(method string, params []interface{}) (interface{}, err
  return nil, err
  }
 
- timeout := time.After(ws.timeout)
+ timeout := time.After(ws.Timeout)
 
  select {
  case <-timeout:
@@ -147,7 +147,7 @@ func (ws *WebSocket) Send(method string, params []interface{}) (interface{}, err
 }
 
 func (ws *WebSocket) read(v interface{}) error {
- _, data, err := ws.conn.ReadMessage()
+ _, data, err := ws.Conn.ReadMessage()
  if err != nil {
  return err
  }
@@ -163,7 +163,7 @@ func (ws *WebSocket) write(v interface{}) error {
 
  ws.connLock.Lock()
  defer ws.connLock.Unlock()
- return ws.conn.WriteMessage(websocket.TextMessage, data)
+ return ws.Conn.WriteMessage(websocket.TextMessage, data)
 }
 
 func (ws *WebSocket) initialize() {