136 lines
3.6 KiB
Go
136 lines
3.6 KiB
Go
package store
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"g.fc.run/ai/byre/internal/api"
|
|
)
|
|
|
|
func migrateWorkflowDefinitions(tx *sql.Tx, _ string) error {
|
|
_, err := tx.Exec(`
|
|
CREATE TABLE workflow_definitions (
|
|
id TEXT PRIMARY KEY,
|
|
description TEXT NOT NULL
|
|
);
|
|
CREATE TABLE workflow_steps (
|
|
definition_id TEXT NOT NULL REFERENCES workflow_definitions(id),
|
|
id INTEGER NOT NULL,
|
|
parent INTEGER,
|
|
position INTEGER NOT NULL,
|
|
kind TEXT NOT NULL CHECK (kind IN ('stage', 'order', 'parallel')),
|
|
stage TEXT NOT NULL,
|
|
model TEXT NOT NULL,
|
|
prompt TEXT NOT NULL,
|
|
PRIMARY KEY (definition_id, id),
|
|
UNIQUE (definition_id, parent, position)
|
|
);
|
|
CREATE TABLE workflow_agents (
|
|
definition_id TEXT NOT NULL REFERENCES workflow_definitions(id),
|
|
step_id INTEGER NOT NULL,
|
|
position INTEGER NOT NULL,
|
|
name TEXT NOT NULL,
|
|
model TEXT NOT NULL,
|
|
prompt TEXT NOT NULL,
|
|
PRIMARY KEY (definition_id, step_id, position),
|
|
FOREIGN KEY (definition_id, step_id) REFERENCES workflow_steps(definition_id, id)
|
|
);
|
|
ALTER TABLE tasks ADD COLUMN workflow_definition_id TEXT REFERENCES workflow_definitions(id);
|
|
ALTER TABLE reviews ADD COLUMN definition_id TEXT REFERENCES workflow_definitions(id);
|
|
ALTER TABLE workflows ADD COLUMN definition_id TEXT REFERENCES workflow_definitions(id);
|
|
`)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
ids := map[string]string{}
|
|
for _, c := range []struct{ table, key, from, to string }{
|
|
{"workflows", "name", "definition", "definition_id"},
|
|
{"tasks", "id", "workflow_definition", "workflow_definition_id"},
|
|
{"reviews", "id", "definition", "definition_id"},
|
|
} {
|
|
err = moveDefinitions(tx, ids, c.table, c.key, c.from, c.to)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
_, err = tx.Exec(`
|
|
CREATE TABLE workflows_next (
|
|
name TEXT PRIMARY KEY,
|
|
definition_id TEXT NOT NULL REFERENCES workflow_definitions(id)
|
|
);
|
|
INSERT INTO workflows_next (name, definition_id) SELECT name, definition_id FROM workflows;
|
|
DROP TABLE workflows;
|
|
ALTER TABLE workflows_next RENAME TO workflows;
|
|
ALTER TABLE tasks DROP COLUMN workflow_definition;
|
|
ALTER TABLE reviews DROP COLUMN definition;
|
|
`)
|
|
return err
|
|
}
|
|
|
|
func moveDefinitions(tx *sql.Tx, ids map[string]string, table, key, from, to string) error {
|
|
rows, err := tx.Query(fmt.Sprintf(`SELECT %s, %s FROM %s WHERE %s != ''`, key, from, table, from))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
type row struct{ key, def string }
|
|
all := []row{}
|
|
for rows.Next() {
|
|
r := row{}
|
|
err = rows.Scan(&r.key, &r.def)
|
|
if err != nil {
|
|
rows.Close()
|
|
return err
|
|
}
|
|
all = append(all, r)
|
|
}
|
|
rows.Close()
|
|
err = rows.Err()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
ctx := context.Background()
|
|
for _, r := range all {
|
|
id, found := ids[r.def]
|
|
if !found {
|
|
def := api.WorkflowDefinition{}
|
|
dec := json.NewDecoder(bytes.NewReader([]byte(r.def)))
|
|
dec.DisallowUnknownFields()
|
|
err = dec.Decode(&def)
|
|
if err != nil {
|
|
return fmt.Errorf("%s %s: %w", table, r.key, err)
|
|
}
|
|
id, err = insertDefinition(ctx, tx, def)
|
|
if err != nil {
|
|
return fmt.Errorf("%s %s: %w", table, r.key, err)
|
|
}
|
|
back, err := loadDefinition(ctx, tx, id)
|
|
if err != nil {
|
|
return fmt.Errorf("%s %s: %w", table, r.key, err)
|
|
}
|
|
want, err := json.Marshal(def)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
got, err := json.Marshal(back)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !bytes.Equal(want, got) {
|
|
return fmt.Errorf("%s %s: definition reads back as %s, not %s", table, r.key, got, want)
|
|
}
|
|
ids[r.def] = id
|
|
}
|
|
_, err = tx.Exec(fmt.Sprintf(`UPDATE %s SET %s = ? WHERE %s = ?`, table, to, key), id, r.key)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|