fetch_ml/internal/api/ws_datasets.go

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))
}