slask

pkg/datastore/datastore.go

@@ -35,52 +35,124 @@ func (ds *DataStore) Close() error {
 	return ds.db.Close()
 }
 
-// initTables creates the necessary database tables
+// initTables creates the necessary database tables and applies migrations
 func (ds *DataStore) initTables() error {
-	tables := []string{
-		`CREATE TABLE IF NOT EXISTS devices (
-			id INTEGER PRIMARY KEY,
-			name TEXT,
-			unique_id TEXT UNIQUE
-		)`,
-		`CREATE TABLE IF NOT EXISTS sensors (
-			sensor_id INTEGER PRIMARY KEY AUTOINCREMENT,
-			protocol TEXT,
-			model TEXT,
-			id INTEGER,
-			name TEXT,
-			temperature_unique_id TEXT,
-			humidity_unique_id TEXT,
-			last_temperature TEXT,
-			last_humidity TEXT,
-			last_timestamp INTEGER,
-			hidden INTEGER DEFAULT 0,
-			UNIQUE(protocol, model, id)
-		)`,
-		`CREATE TABLE IF NOT EXISTS potential_devices (
-			id INTEGER PRIMARY KEY AUTOINCREMENT,
-			class TEXT,
-			protocol TEXT,
-			model TEXT,
-			device_id TEXT,
-			last_data TEXT,
-			last_seen INTEGER
-		)`,
+	// Create schema version table
+	if _, err := ds.db.Exec(`
+		CREATE TABLE IF NOT EXISTS schema_version (
+			version INTEGER PRIMARY KEY,
+			applied_at INTEGER NOT NULL
+		)
+	`); err != nil {
+		return err
 	}
 
-	for _, table := range tables {
-		if _, err := ds.db.Exec(table); err != nil {
-			return err
+	// Get current schema version
+	currentVersion := ds.getCurrentSchemaVersion()
+
+	// Apply migrations
+	migrations := []struct {
+		version int
+		sql     []string
+	}{
+		{
+			version: 1,
+			sql: []string{
+				`CREATE TABLE IF NOT EXISTS devices (
+					id INTEGER PRIMARY KEY,
+					name TEXT,
+					unique_id TEXT UNIQUE
+				)`,
+				`CREATE TABLE IF NOT EXISTS sensors (
+					sensor_id INTEGER PRIMARY KEY AUTOINCREMENT,
+					protocol TEXT,
+					model TEXT,
+					id INTEGER,
+					name TEXT,
+					temperature_unique_id TEXT,
+					humidity_unique_id TEXT,
+					last_temperature TEXT,
+					last_humidity TEXT,
+					last_timestamp INTEGER,
+					hidden INTEGER DEFAULT 0,
+					UNIQUE(protocol, model, id)
+				)`,
+				`CREATE TABLE IF NOT EXISTS potential_devices (
+					id INTEGER PRIMARY KEY AUTOINCREMENT,
+					class TEXT,
+					protocol TEXT,
+					model TEXT,
+					device_id TEXT,
+					last_data TEXT,
+					last_seen INTEGER
+				)`,
+			},
+		},
+		{
+			version: 2,
+			sql: []string{
+				`ALTER TABLE devices ADD COLUMN protocol TEXT`,
+				`ALTER TABLE devices ADD COLUMN model TEXT`,
+			},
+		},
+	}
+
+	// Apply pending migrations
+	for _, migration := range migrations {
+		if migration.version > currentVersion {
+			for _, sql := range migration.sql {
+				if _, err := ds.db.Exec(sql); err != nil {
+					// Ignore errors for ALTER TABLE ADD COLUMN if column already exists
+					if !ds.isColumnExistsError(err) {
+						return fmt.Errorf("migration v%d failed: %w", migration.version, err)
+					}
+				}
+			}
+
+			// Record migration
+			if err := ds.recordMigration(migration.version); err != nil {
+				return err
+			}
 		}
 	}
+
 	return nil
 }
 
+// getCurrentSchemaVersion returns the current schema version
+func (ds *DataStore) getCurrentSchemaVersion() int {
+	var version int
+	err := ds.db.QueryRow("SELECT MAX(version) FROM schema_version").Scan(&version)
+	if err != nil {
+		return 0
+	}
+	return version
+}
+
+// recordMigration records that a migration has been applied
+func (ds *DataStore) recordMigration(version int) error {
+	_, err := ds.db.Exec(
+		"INSERT INTO schema_version (version, applied_at) VALUES (?, ?)",
+		version, time.Now().Unix(),
+	)
+	return err
+}
+
+// isColumnExistsError checks if the error is due to a column already existing
+func (ds *DataStore) isColumnExistsError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return sql.ErrNoRows != err &&
+		(err.Error() == "duplicate column name: protocol" ||
+			err.Error() == "duplicate column name: model")
+}
+
 // UpsertDevice inserts or updates a device
 func (ds *DataStore) UpsertDevice(device *Device) error {
 	_, err := ds.db.Exec(
-		"INSERT OR REPLACE INTO devices (id, name, unique_id) VALUES (?, ?, ?)",
-		device.ID, device.Name, device.UniqueID,
+		"INSERT OR REPLACE INTO devices (id, name, unique_id, protocol, model) VALUES (?, ?, ?, ?, ?)",
+		device.ID, device.Name, device.UniqueID, device.Protocol, device.Model,
 	)
 	return err
 }
@@ -89,9 +161,9 @@ func (ds *DataStore) UpsertDevice(device *Device) error {
 func (ds *DataStore) GetDevice(id int) (*Device, error) {
 	device := &Device{}
 	err := ds.db.QueryRow(
-		"SELECT id, name, unique_id FROM devices WHERE id = ?",
+		"SELECT id, name, unique_id, protocol, model FROM devices WHERE id = ?",
 		id,
-	).Scan(&device.ID, &device.Name, &device.UniqueID)
+	).Scan(&device.ID, &device.Name, &device.UniqueID, &device.Protocol, &device.Model)
 	if err != nil {
 		return nil, err
 	}
@@ -101,7 +173,7 @@ func (ds *DataStore) GetDevice(id int) (*Device, error) {
 // ListDevices returns an iterator over all devices
 func (ds *DataStore) ListDevices() iter.Seq[*Device] {
 	return func(yield func(*Device) bool) {
-		rows, err := ds.db.Query("SELECT id, name, unique_id FROM devices ORDER BY id")
+		rows, err := ds.db.Query("SELECT id, name, unique_id, protocol, model FROM devices ORDER BY id")
 		if err != nil {
 			return
 		}
@@ -109,7 +181,7 @@ func (ds *DataStore) ListDevices() iter.Seq[*Device] {
 
 		for rows.Next() {
 			device := &Device{}
-			if err := rows.Scan(&device.ID, &device.Name, &device.UniqueID); err != nil {
+			if err := rows.Scan(&device.ID, &device.Name, &device.UniqueID, &device.Protocol, &device.Model); err != nil {
 				continue
 			}
 			if !yield(device) {