// Package datasets provides WebSocket handlers for dataset-related operations package datasets import ( "context" "encoding/binary" "encoding/json" "time" "github.com/gorilla/websocket" "github.com/jfraeys/fetch_ml/internal/auth" "github.com/jfraeys/fetch_ml/internal/logging" "github.com/jfraeys/fetch_ml/internal/storage" ) // Handler provides dataset-related WebSocket handlers type Handler struct { logger *logging.Logger db *storage.DB dataDir string } // NewHandler creates a new datasets handler func NewHandler( logger *logging.Logger, db *storage.DB, dataDir string, ) *Handler { return &Handler{ logger: logger, db: db, dataDir: dataDir, } } // Error codes const ( ErrorCodeInvalidRequest = 0x01 ErrorCodeAuthenticationFailed = 0x02 ErrorCodePermissionDenied = 0x03 ErrorCodeResourceNotFound = 0x04 ) // sendErrorPacket sends an error response packet to the client func (h *Handler) sendErrorPacket(conn *websocket.Conn, code byte, message, details string) error { err := map[string]interface{}{ "error": true, "code": code, "message": message, "details": details, } return conn.WriteJSON(err) } // sendSuccessPacket sends a success response packet func (h *Handler) sendSuccessPacket(conn *websocket.Conn, data map[string]interface{}) error { return conn.WriteJSON(data) } // sendDataPacket sends a data response packet func (h *Handler) sendDataPacket(conn *websocket.Conn, dataType string, payload []byte) error { return conn.WriteJSON(map[string]interface{}{ "type": dataType, "payload": string(payload), }) } // HandleDatasetList handles listing datasets // Protocol: [api_key_hash:16] func (h *Handler) HandleDatasetList(conn *websocket.Conn, payload []byte, user *auth.User) error { h.logger.Info("listing datasets", "user", user.Name) var datasets []*storage.Dataset if h.db != nil { var err error datasets, err = h.db.ListDatasets(context.Background(), 100) if err != nil { h.logger.Warn("failed to list datasets from db", "error", err) datasets = []*storage.Dataset{} } } data, _ := json.Marshal(datasets) return h.sendDataPacket(conn, "datasets", data) } // HandleDatasetRegister handles registering a new dataset // Protocol: [api_key_hash:16][name_len:1][name:var][path_len:2][path:var] func (h *Handler) HandleDatasetRegister(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "register dataset payload too short", "") } offset := 16 nameLen := int(payload[offset]) offset++ if nameLen <= 0 || len(payload) < offset+nameLen+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid name length", "") } name := string(payload[offset : offset+nameLen]) offset += nameLen pathLen := int(binary.BigEndian.Uint16(payload[offset : offset+2])) offset += 2 if pathLen < 0 || len(payload) < offset+pathLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid path length", "") } path := string(payload[offset : offset+pathLen]) h.logger.Info("registering dataset", "name", name, "path", path, "user", user.Name) // Save to database if available if h.db != nil { ds := &storage.Dataset{ Name: name, URL: path, } if err := h.db.UpsertDataset(context.Background(), ds); err != nil { h.logger.Warn("failed to save dataset to db", "error", err, "name", name) } } return h.sendSuccessPacket(conn, map[string]interface{}{ "success": true, "name": name, "path": path, "user": user.Name, "time": time.Now().UTC(), }) } // HandleDatasetInfo handles getting dataset info // Protocol: [api_key_hash:16][dataset_id_len:1][dataset_id:var] func (h *Handler) HandleDatasetInfo(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "dataset info payload too short", "") } offset := 16 datasetIDLen := int(payload[offset]) offset++ if datasetIDLen <= 0 || len(payload) < offset+datasetIDLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid dataset ID length", "") } datasetID := string(payload[offset : offset+datasetIDLen]) h.logger.Info("getting dataset info", "dataset_id", datasetID, "user", user.Name) // Query database if available if h.db != nil { ds, err := h.db.GetDataset(context.Background(), datasetID) if err == nil && ds != nil { data, _ := json.Marshal(ds) return h.sendDataPacket(conn, "dataset_info", data) } if err != nil { h.logger.Warn("failed to get dataset from db", "error", err, "name", datasetID) } } return h.sendDataPacket(conn, "dataset_info", []byte("{}")) } // HandleDatasetSearch handles searching datasets // Protocol: [api_key_hash:16][query_len:2][query:var] func (h *Handler) HandleDatasetSearch(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "dataset search payload too short", "") } offset := 16 queryLen := int(binary.BigEndian.Uint16(payload[offset : offset+2])) offset += 2 if queryLen < 0 || len(payload) < offset+queryLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid query length", "") } query := string(payload[offset : offset+queryLen]) h.logger.Info("searching datasets", "query", query, "user", user.Name) // Search database if available var datasets []*storage.Dataset if h.db != nil { var err error datasets, err = h.db.SearchDatasets(context.Background(), query, 100) if err != nil { h.logger.Warn("failed to search datasets in db", "error", err, "query", query) datasets = []*storage.Dataset{} } } data, _ := json.Marshal(datasets) return h.sendDataPacket(conn, "datasets", data) }