add more handlers and some comments

engine.go

@@ -4,6 +4,7 @@ import (
 	"log/slog"
 	"net/http"
 
+	"github.com/gorilla/websocket"
 	"github.com/wisplite/tether/reactivity"
 	"gorm.io/gorm"
 )
@@ -21,12 +22,12 @@ func NewEngine(db *gorm.DB) *Engine {
 }
 
 func (e *Engine) RegisterMutation(name string, mutation func(ctx *MutationCtx) error) {
-	e.mutations[name] = mutation
+	e.mutations[name] = mutation // stores the mutation in the list of valid mutations
 	slog.Debug("Registered mutation", "name", name)
 }
 
 func (e *Engine) RegisterQuery(name string, query func(ctx *QueryCtx) error) {
-	e.queries[name] = query
+	e.queries[name] = query // stores the query in the list of valid queries
 	slog.Debug("Registered query", "name", name)
 }
 
@@ -36,10 +37,23 @@ func (e *Engine) CreateTable(name string, schema interface{}) {
 }
 
 func (e *Engine) Handle(w http.ResponseWriter, r *http.Request) {
-	reactivity.Handle(w, r, e)
+	reactivity.Handle(w, r, e) // wraps the raw websocket connection with the engine handler
+}
+
+func (e *Engine) OnConnect(conn *websocket.Conn) error {
+	slog.Debug("Connected to websocket")
+	// TODO: implement the logic to handle the connection
+	return nil
+}
+
+func (e *Engine) OnDisconnect(conn *websocket.Conn) error {
+	slog.Debug("Disconnected from websocket")
+	// TODO: implement the logic to handle the disconnection
+	return nil
 }
 
 func (e *Engine) OnReceiveMessage(msg map[string]interface{}) error {
 	slog.Debug("Received message", "message", msg)
+	// TODO: implement the logic to handle the message
 	return nil
 }

reactivity/websocket.go

@@ -9,13 +9,15 @@ import (
 )
 
 // MessageReceiver receives decoded WebSocket payloads. Implemented by tether.Engine.
-type MessageReceiver interface {
+type EngineHandler interface {
+	OnConnect(conn *websocket.Conn) error
+	OnDisconnect(conn *websocket.Conn) error
 	OnReceiveMessage(msg map[string]interface{}) error
 }
 
 var upgrader = websocket.Upgrader{}
 
-func Handle(w http.ResponseWriter, r *http.Request, e MessageReceiver) {
+func Handle(w http.ResponseWriter, r *http.Request, e EngineHandler) {
 	ws, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		slog.Error("WS: Failed to upgrade to websocket", "error", err)
@@ -23,6 +25,13 @@ func Handle(w http.ResponseWriter, r *http.Request, e MessageReceiver) {
 	}
 	defer ws.Close()
 
+	err = e.OnConnect(ws)
+	if err != nil {
+		slog.Error("WS: Failed to call onConnect handler", "error", err)
+		return
+	}
+	defer e.OnDisconnect(ws)
+
 	for {
 		_, message, err := ws.ReadMessage()
 		if err != nil {