new-server/internal/mysql/devicerepo.go

package mysql

import (
	"context"
	"database/sql"
	"encoding/json"
	"git.aiterp.net/lucifer/new-server/models"
	sq "github.com/Masterminds/squirrel"
	"github.com/jmoiron/sqlx"
	"log"
	"strings"
)

type deviceRecord struct {
	ID                  int             `db:"id"`
	BridgeID            int             `db:"bridge_id"`
	InternalID          string          `db:"internal_id"`
	Icon                string          `db:"icon"`
	Name                string          `db:"name"`
	Capabilities        string          `db:"capabilities"`
	ButtonNames         string          `db:"button_names"`
	SceneAssignmentJSON json.RawMessage `db:"scene_assignments"`
}

type deviceStateRecord struct {
	DeviceID   int     `db:"device_id"`
	Hue        float64 `db:"hue"`
	Saturation float64 `db:"saturation"`
	Kelvin     int     `db:"kelvin"`
	Power      bool    `db:"power"`
	Intensity  float64 `db:"intensity"`
}

type devicePropertyRecord struct {
	DeviceID int    `db:"device_id"`
	Key      string `db:"prop_key"`
	Value    string `db:"prop_value"`
	IsUser   bool   `db:"is_user"`
}

type deviceTagRecord struct {
	DeviceID int    `db:"device_id"`
	TagName  string `db:"tag_name"`
}

type DeviceRepo struct {
	DBX *sqlx.DB
}

func (r *DeviceRepo) Find(ctx context.Context, id int) (*models.Device, error) {
	var device deviceRecord
	err := r.DBX.GetContext(ctx, &device, "SELECT * FROM device WHERE id = ?", id)
	if err != nil {
		return nil, dbErr(err)
	}

	return r.populateOne(ctx, device)
}

func (r *DeviceRepo) FetchByReference(ctx context.Context, kind models.ReferenceKind, value string) ([]models.Device, error) {
	q := sq.Select("device.*").From("device")

	var requiredTags []string
	var unwantedTags []string

	switch kind {
	case models.RKDeviceID:
		q = q.Where(sq.Eq{"id": strings.Split(value, ",")})
	case models.RKBridgeID:
		q = q.Where(sq.Eq{"bridge_id": strings.Split(value, ",")})
	case models.RKName:
		value = strings.ReplaceAll(value, "*", "%")
		q = q.Where(sq.Like{"name": value})
	case models.RKTag:
		allTags := strings.Split(strings.ReplaceAll(strings.ReplaceAll(value, "-", ",-"), "+", ",+"), ",")
		optionalTags := make([]string, 0, len(allTags))

		for _, tag := range allTags {
			if tag == "" || tag == "+" || tag == "-" {
				continue
			}

			if strings.HasPrefix(tag, "+") {
				requiredTags = append(requiredTags, tag[1:])
			} else if strings.HasPrefix(tag, "-") {
				unwantedTags = append(unwantedTags, tag[1:])
			} else {
				optionalTags = append(optionalTags, tag)
			}
		}

		q = q.Join("device_tag dt ON device.id=dt.device_id").Where(sq.Eq{"dt.tag_name": optionalTags})
	case models.RKAll:
	default:
		log.Println("Unknown reference kind used for device fetch:", kind)
		return []models.Device{}, nil
	}

	query, args, err := q.OrderBy("name", "id").ToSql()
	if err != nil {
		if err == sql.ErrNoRows {
			return []models.Device{}, nil
		}

		return nil, dbErr(err)
	}

	records := make([]deviceRecord, 0, 8)
	err = r.DBX.SelectContext(ctx, &records, query, args...)
	if err != nil {
		return nil, dbErr(err)
	}

	if len(requiredTags) > 0 || len(unwantedTags) > 0 {
		return r.populateFiltered(ctx, records, requiredTags, unwantedTags)
	}

	return r.populate(ctx, records)
}

func (r *DeviceRepo) SaveMany(ctx context.Context, mode models.SaveMode, devices []models.Device) error {
	tx, err := r.DBX.Beginx()
	if err != nil {
		return dbErr(err)
	}
	defer tx.Rollback()

	for i, device := range devices {
		scenesJSON, err := json.Marshal(device.SceneAssignments)
		if err != nil {
			return err
		}

		record := deviceRecord{
			ID:                  device.ID,
			BridgeID:            device.BridgeID,
			InternalID:          device.InternalID,
			SceneAssignmentJSON: scenesJSON,
			Icon:                device.Icon,
			Name:                device.Name,
			Capabilities:        strings.Join(models.DeviceCapabilitiesToStrings(device.Capabilities), ","),
			ButtonNames:         strings.Join(device.ButtonNames, ","),
		}

		if device.ID > 0 {
			_, err := tx.NamedExecContext(ctx, `
			UPDATE device SET
				internal_id = :internal_id,
				icon = :icon,
				name = :name,
				capabilities = :capabilities,
				button_names = :button_names,
 				scene_assignments = :scene_assignments
			WHERE id=:id
		`, record)
			if err != nil {
				return dbErr(err)
			}

			// Let's just be lazy for now, optimize later if need be.
			if mode == 0 || mode&models.SMTags != 0 {
				_, err = tx.ExecContext(ctx, "DELETE FROM device_tag WHERE device_id=?", record.ID)
				if err != nil {
					return dbErr(err)
				}
			}

			if mode == 0 || mode&models.SMProperties != 0 {
				_, err = tx.ExecContext(ctx, "DELETE FROM device_property WHERE device_id=?", record.ID)
				if err != nil {
					return dbErr(err)
				}
			}
		} else {
			res, err := tx.NamedExecContext(ctx, `
			INSERT INTO device (bridge_id, internal_id, icon, name, capabilities, button_names)
			VALUES (:bridge_id, :internal_id, :icon, :name, :capabilities, :button_names)
		`, record)
			if err != nil {
				return dbErr(err)
			}

			lastID, err := res.LastInsertId()
			if err != nil {
				return dbErr(err)
			}

			record.ID = int(lastID)
			devices[i].ID = int(lastID)
		}

		if mode == 0 || mode&models.SMTags != 0 {
			for _, tag := range device.Tags {
				_, err := tx.ExecContext(ctx, "INSERT INTO device_tag (device_id, tag_name) VALUES (?, ?)", record.ID, tag)
				if err != nil {
					return dbErr(err)
				}
			}
		}

		if mode == 0 || mode&models.SMProperties != 0 {
			for key, value := range device.UserProperties {
				_, err := tx.ExecContext(ctx, "INSERT INTO device_property (device_id, prop_key, prop_value, is_user) VALUES (?, ?, ?, 1)",
					record.ID, key, value,
				)
				if err != nil {
					return dbErr(err)
				}
			}

			for key, value := range device.DriverProperties {
				j, err := json.Marshal(value)
				if err != nil {
					// Eh, it'll get filled by the driver anyway
					continue
				}

				_, err = tx.ExecContext(ctx, "INSERT INTO device_property (device_id, prop_key, prop_value, is_user) VALUES (?, ?, ?, 0)",
					record.ID, key, string(j),
				)
				if err != nil {
					// Return err here anyway, it might put the tx in a bad state to ignore it.
					return dbErr(err)
				}
			}
		}

		if mode == 0 || mode&models.SMState != 0 {
			_, err = tx.NamedExecContext(ctx, `
			REPLACE INTO device_state(device_id, hue, saturation, kelvin, power, intensity) 
			VALUES (:device_id, :hue, :saturation, :kelvin, :power, :intensity)
		`, deviceStateRecord{
				DeviceID:   record.ID,
				Hue:        device.State.Color.Hue,
				Saturation: device.State.Color.Saturation,
				Kelvin:     device.State.Color.Kelvin,
				Power:      device.State.Power,
				Intensity:  device.State.Intensity,
			})
			if err != nil {
				return dbErr(err)
			}
		}
	}

	return tx.Commit()
}

func (r *DeviceRepo) Save(ctx context.Context, device *models.Device, mode models.SaveMode) error {
	devices := []models.Device{*device}

	err := r.SaveMany(ctx, mode, devices)
	if err != nil {
		return err
	}

	*device = devices[0]
	return nil
}

func (r *DeviceRepo) Delete(ctx context.Context, device *models.Device) error {
	_, err := r.DBX.ExecContext(ctx, "DELETE FROM device WHERE Id=?", device.ID)
	if err != nil {
		return dbErr(err)
	}

	return nil
}

func (r *DeviceRepo) populateOne(ctx context.Context, record deviceRecord) (*models.Device, error) {
	records, err := r.populate(ctx, []deviceRecord{record})
	if err != nil {
		return nil, err
	}

	return &records[0], nil
}

func (r *DeviceRepo) populateFiltered(ctx context.Context, records []deviceRecord, requiredTags []string, unwantedTags []string) ([]models.Device, error) {
	devices, err := r.populate(ctx, records)
	if err != nil {
		return nil, err
	}

	filteredDevices := make([]models.Device, 0, len(devices))

	for _, device := range devices {
		if device.HasTag(unwantedTags...) {
			continue
		}

		if len(requiredTags) == 0 || device.HasTag(requiredTags...) {
			filteredDevices = append(filteredDevices, device)
		}
	}

	return filteredDevices, nil
}

func (r *DeviceRepo) populate(ctx context.Context, records []deviceRecord) ([]models.Device, error) {
	if len(records) == 0 {
		return []models.Device{}, nil
	}

	ids := make([]int, 0, len(records))
	for _, record := range records {
		ids = append(ids, record.ID)
	}

	tagsQuery, tagsArgs, err := sq.Select("*").From("device_tag").Where(sq.Eq{"device_id": ids}).ToSql()
	if err != nil {
		return nil, dbErr(err)
	}
	propsQuery, propsArgs, err := sq.Select("*").From("device_property").Where(sq.Eq{"device_id": ids}).ToSql()
	if err != nil {
		return nil, dbErr(err)
	}
	stateQuery, stateArgs, err := sq.Select("*").From("device_state").Where(sq.Eq{"device_id": ids}).ToSql()
	if err != nil {
		return nil, dbErr(err)
	}

	states := make([]deviceStateRecord, 0, len(records))
	props := make([]devicePropertyRecord, 0, len(records)*8)
	tags := make([]deviceTagRecord, 0, len(records)*4)

	err = r.DBX.SelectContext(ctx, &states, stateQuery, stateArgs...)
	if err != nil {
		return nil, dbErr(err)
	}
	err = r.DBX.SelectContext(ctx, &props, propsQuery, propsArgs...)
	if err != nil {
		return nil, dbErr(err)
	}
	err = r.DBX.SelectContext(ctx, &tags, tagsQuery, tagsArgs...)
	if err != nil {
		return nil, dbErr(err)
	}

	hasAdded := make(map[int]bool, len(records))
	devices := make([]models.Device, 0, len(records))
	for _, record := range records {
		if hasAdded[record.ID] {
			continue
		}

		device := models.Device{
			ID:               record.ID,
			BridgeID:         record.BridgeID,
			InternalID:       record.InternalID,
			Icon:             record.Icon,
			Name:             record.Name,
			SceneAssignments: make([]models.DeviceSceneAssignment, 0, 4),
			ButtonNames:      strings.Split(record.ButtonNames, ","),
			DriverProperties: make(map[string]interface{}, 8),
			UserProperties:   make(map[string]string, 8),
			Tags:             make([]string, 0, 8),
		}

		_ = json.Unmarshal(record.SceneAssignmentJSON, &device.SceneAssignments)

		if device.ButtonNames[0] == "" {
			device.ButtonNames = device.ButtonNames[:0]
		}

		caps := make([]models.DeviceCapability, 0, 16)
		for _, capStr := range strings.Split(record.Capabilities, ",") {
			caps = append(caps, models.DeviceCapability(capStr))
		}
		device.Capabilities = caps

		for _, state := range states {
			if state.DeviceID == record.ID {
				device.State = models.DeviceState{
					Power: state.Power,
					Color: models.ColorValue{
						Hue:        state.Hue,
						Saturation: state.Saturation,
						Kelvin:     state.Kelvin,
					},
					Intensity: state.Intensity,
				}
			}
		}

		driverProps := make(map[string]json.RawMessage, 8)
		for _, prop := range props {
			if prop.DeviceID == record.ID {
				if prop.IsUser {
					device.UserProperties[prop.Key] = prop.Value
				} else {
					driverProps[prop.Key] = json.RawMessage(prop.Value)
				}
			}
		}
		if len(driverProps) > 0 {
			j, err := json.Marshal(driverProps)
			if err != nil {
				return nil, dbErr(err)
			}
			err = json.Unmarshal(j, &device.DriverProperties)
			if err != nil {
				return nil, dbErr(err)
			}
		}

		for _, tag := range tags {
			if tag.DeviceID == record.ID {
				device.Tags = append(device.Tags, tag.TagName)
			}
		}

		hasAdded[record.ID] = true
		devices = append(devices, device)
	}

	return devices, nil
}