Fix #124

db.go

@@ -37,8 +37,8 @@ func New(url string, connection conn.Connection) (*DB, error) {
 // --------------------------------------------------
 
 // Close closes the underlying WebSocket connection.
-func (db *DB) Close() {
- _ = db.conn.Close()
+func (db *DB) Close() error {
+ return db.conn.Close()
 }
 
 // --------------------------------------------------

db_test.go

@@ -4,6 +4,7 @@ import (
  "bytes"
  "encoding/json"
  "fmt"
+ "io"
  rawslog "log/slog"
  "os"
  "sync"
@@ -31,6 +32,7 @@ type SurrealDBTestSuite struct {
  db *surrealdb.DB
  name string
  connImplementations map[string]conn.Connection
+ logBuffer *bytes.Buffer
 }
 
 // a simple user struct for testing
@@ -55,20 +57,23 @@ func TestSurrealDBSuite(t *testing.T) {
  SurrealDBSuite.connImplementations = make(map[string]conn.Connection)
 
  // Without options
- logData := createLogger(t)
+ buff := bytes.NewBufferString("")
+ logData := createLogger(t, buff)
  SurrealDBSuite.connImplementations["gorilla"] = gorilla.Create().Logger(logData)
+ SurrealDBSuite.logBuffer = buff
 
  // With options
- logData = createLogger(t)
- SurrealDBSuite.connImplementations["gorilla_opt"] = gorilla.Create().SetTimeOut(time.Minute).SetCompression(true).Logger(logData)
+ buffOpt := bytes.NewBufferString("")
+ logDataOpt := createLogger(t, buff)
+ SurrealDBSuite.connImplementations["gorilla_opt"] = gorilla.Create().SetTimeOut(time.Minute).SetCompression(true).Logger(logDataOpt)
+ SurrealDBSuite.logBuffer = buffOpt
 
  RunWsMap(t, SurrealDBSuite)
 }
 
-func createLogger(t *testing.T) logger.Logger {
+func createLogger(t *testing.T, writer io.Writer) logger.Logger {
  t.Helper()
- buff := bytes.NewBuffer([]byte{})
- handler := rawslog.NewJSONHandler(buff, &rawslog.HandlerOptions{Level: rawslog.LevelDebug})
+ handler := rawslog.NewJSONHandler(writer, &rawslog.HandlerOptions{Level: rawslog.LevelDebug})
  return slog.New(handler)
 }
 
@@ -86,11 +91,16 @@ func RunWsMap(t *testing.T, s *SurrealDBTestSuite) {
 func (s *SurrealDBTestSuite) TearDownTest() {
  _, err := s.db.Delete("users")
  s.Require().NoError(err)
+
+ if s.logBuffer.Len() > 0 {
+ s.T().Logf("Log output:\n%s", s.logBuffer.String())
+ }
 }
 
 // TearDownSuite is called after the s has finished running
 func (s *SurrealDBTestSuite) TearDownSuite() {
- s.db.Close()
+ err := s.db.Close()
+ s.Require().NoError(err)
 }
 
 func (t testUser) String() (str string, err error) {

pkg/conn/gorilla/gorilla.go

@@ -4,6 +4,7 @@ import (
  "encoding/json"
  "errors"
  "fmt"
+ "net"
  "reflect"
  "strconv"
  "sync"
@@ -104,11 +105,15 @@ func (ws *WebSocket) SetCompression(compress bool) *WebSocket {
 }
 
 func (ws *WebSocket) Close() error {
- defer func() {
- close(ws.close)
- }()
+ ws.connLock.Lock()
+ defer ws.connLock.Unlock()
+ close(ws.close)
+ err := ws.Conn.WriteMessage(gorilla.CloseMessage, gorilla.FormatCloseMessage(CloseMessageCode, ""))
+ if err != nil {
+ return err
+ }
 
- return ws.Conn.WriteMessage(gorilla.CloseMessage, gorilla.FormatCloseMessage(CloseMessageCode, ""))
+ return ws.Conn.Close()
 }
 
 func (ws *WebSocket) LiveNotifications(liveQueryID string) (chan model.Notification, error) {
@@ -239,6 +244,9 @@ func (ws *WebSocket) initialize() {
  var res rpc.RPCResponse
  err := ws.read(&res)
  if err != nil {
+ if errors.Is(err, net.ErrClosed) {
+ break
+ }
  ws.logger.Error(err.Error())
  continue
  }