package api import ( "database/sql" "encoding/binary" "encoding/json" "net/url" "strings" "github.com/gorilla/websocket" "github.com/jfraeys/fetch_ml/internal/api/helpers" "github.com/jfraeys/fetch_ml/internal/storage" ) func (h *WSHandler) handleDatasetList(conn *websocket.Conn, payload []byte) error { user, err := h.authenticate(conn, payload, ProtocolMinDatasetList) if err != nil { return err } if err := h.requirePermission(user, PermDatasetsRead, conn); err != nil { return err } if err := h.requireDB(conn); err != nil { return err } ctx, cancel := helpers.DBContextShort() defer cancel() datasets, err := h.db.ListDatasets(ctx, 0) if err != nil { return h.sendErrorPacket(conn, ErrorCodeDatabaseError, "Failed to list datasets", err.Error()) } data, err := json.Marshal(datasets) if err != nil { return h.sendErrorPacket( conn, ErrorCodeServerOverloaded, "Failed to serialize response", err.Error(), ) } return h.sendResponsePacket(conn, NewDataPacket("datasets", data)) } func (h *WSHandler) handleDatasetRegister(conn *websocket.Conn, payload []byte) error { user, err := h.authenticate(conn, payload, ProtocolMinDatasetRegister) if err != nil { return err } if err := h.requirePermission(user, PermDatasetsCreate, conn); err != nil { return err } if err := h.requireDB(conn); err != nil { return err } offset := ProtocolAPIKeyHashLen nameLen := int(payload[offset]) offset++ if nameLen <= 0 || len(payload) < offset+nameLen+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid dataset name length", "") } name := string(payload[offset : offset+nameLen]) offset += nameLen urlLen := int(binary.BigEndian.Uint16(payload[offset : offset+2])) offset += 2 if urlLen <= 0 || len(payload) < offset+urlLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid dataset url length", "") } urlStr := string(payload[offset : offset+urlLen]) if strings.TrimSpace(name) == "" { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "dataset name required", "") } if u, err := url.Parse(urlStr); err != nil || u.Scheme == "" { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid dataset url", "") } ctx, cancel := helpers.DBContextShort() defer cancel() if err := h.db.UpsertDataset(ctx, &storage.Dataset{Name: name, URL: urlStr}); err != nil { return h.sendErrorPacket(conn, ErrorCodeDatabaseError, "Failed to register dataset", err.Error()) } return h.sendResponsePacket(conn, NewSuccessPacket("Dataset registered")) } func (h *WSHandler) handleDatasetInfo(conn *websocket.Conn, payload []byte) error { user, err := h.authenticate(conn, payload, ProtocolMinDatasetInfo) if err != nil { return err } if err := h.requirePermission(user, PermDatasetsRead, conn); err != nil { return err } if err := h.requireDB(conn); err != nil { return err } offset := ProtocolAPIKeyHashLen nameLen := int(payload[offset]) offset++ if nameLen <= 0 || len(payload) < offset+nameLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid dataset name length", "") } name := string(payload[offset : offset+nameLen]) ctx, cancel := helpers.DBContextShort() defer cancel() ds, err := h.db.GetDataset(ctx, name) if err != nil { if err == sql.ErrNoRows { return h.sendErrorPacket(conn, ErrorCodeResourceNotFound, "Dataset not found", "") } return h.sendErrorPacket(conn, ErrorCodeDatabaseError, "Failed to get dataset", err.Error()) } data, err := json.Marshal(ds) if err != nil { return h.sendErrorPacket( conn, ErrorCodeServerOverloaded, "Failed to serialize response", err.Error(), ) } return h.sendResponsePacket(conn, NewDataPacket("dataset", data)) } func (h *WSHandler) handleDatasetSearch(conn *websocket.Conn, payload []byte) error { user, err := h.authenticate(conn, payload, ProtocolMinDatasetSearch) if err != nil { return err } if err := h.requirePermission(user, PermDatasetsRead, conn); err != nil { return err } if err := h.requireDB(conn); err != nil { return err } offset := ProtocolAPIKeyHashLen termLen := int(payload[offset]) offset++ if termLen < 0 || len(payload) < offset+termLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid search term length", "") } term := string(payload[offset : offset+termLen]) term = strings.TrimSpace(term) ctx, cancel := helpers.DBContextShort() defer cancel() datasets, err := h.db.SearchDatasets(ctx, term, 0) if err != nil { return h.sendErrorPacket(conn, ErrorCodeDatabaseError, "Failed to search datasets", err.Error()) } data, err := json.Marshal(datasets) if err != nil { return h.sendErrorPacket( conn, ErrorCodeServerOverloaded, "Failed to serialize response", err.Error(), ) } return h.sendResponsePacket(conn, NewDataPacket("datasets", data)) }