diff --git a/cmd/main-server.go b/cmd/main-server.go index 1e1a2e5d..6d4488e9 100644 --- a/cmd/main-server.go +++ b/cmd/main-server.go @@ -145,12 +145,6 @@ func HandleGetClientData(w http.ResponseWriter, r *http.Request) { return } cdata = cdata.Clean() - mdata, err := sstore.GetCmdMigrationInfo(r.Context()) - if err != nil { - WriteJsonError(w, err) - return - } - cdata.Migration = mdata WriteJsonSuccess(w, cdata) return } @@ -573,7 +567,6 @@ func main() { time.Sleep(10 * time.Second) pcloud.StartUpdateWriter() }() - go sstore.RunCmdScreenMigration() gr := mux.NewRouter() gr.HandleFunc("/api/ptyout", AuthKeyWrap(HandleGetPtyOut)) gr.HandleFunc("/api/remote-pty", AuthKeyWrap(HandleRemotePty)) diff --git a/pkg/sstore/migrate.go b/pkg/sstore/migrate.go index 93abb367..72b6bc71 100644 --- a/pkg/sstore/migrate.go +++ b/pkg/sstore/migrate.go @@ -8,17 +8,18 @@ import ( "strconv" "time" + sh2db "github.com/commandlinedev/prompt-server/db" _ "github.com/golang-migrate/migrate/v4/database/sqlite3" _ "github.com/golang-migrate/migrate/v4/source/file" "github.com/golang-migrate/migrate/v4/source/iofs" _ "github.com/mattn/go-sqlite3" - sh2db "github.com/commandlinedev/prompt-server/db" "github.com/golang-migrate/migrate/v4" ) const MaxMigration = 19 const MigratePrimaryScreenVersion = 9 +const CmdScreenSpecialMigration = 13 func MakeMigrate() (*migrate.Migrate, error) { fsVar, err := iofs.New(sh2db.MigrationFS, "migrations") @@ -56,46 +57,67 @@ func copyFile(srcFile string, dstFile string) error { return dstFd.Close() } -func MigrateUp() error { +func MigrateUpStep(m *migrate.Migrate, newVersion uint) error { + startTime := time.Now() + err := m.Migrate(newVersion) + if err != nil { + return err + } + if newVersion == CmdScreenSpecialMigration { + mErr := RunCmdScreenMigration13() + if mErr != nil { + return mErr + } + } + log.Printf("[db] migration v%d, elapsed %v\n", newVersion, time.Since(startTime)) + return nil +} + +func MigrateUp(targetVersion uint) error { m, err := MakeMigrate() if err != nil { return err } - curVersion, dirty, err := m.Version() - if err == migrate.ErrNilVersion { - curVersion = 0 - err = nil - } + curVersion, dirty, err := MigrateVersion(m) if dirty { return fmt.Errorf("cannot migrate up, database is dirty") } if err != nil { return fmt.Errorf("cannot get current migration version: %v", err) } - if curVersion >= MaxMigration { + if curVersion >= targetVersion { return nil } - log.Printf("[db] migrating from %d to %d\n", curVersion, MaxMigration) + log.Printf("[db] migrating from %d to %d\n", curVersion, targetVersion) log.Printf("[db] backing up database %s to %s\n", DBFileName, DBFileNameBackup) err = copyFile(GetDBName(), GetDBBackupName()) if err != nil { return fmt.Errorf("error creating database backup: %v", err) } - startTime := time.Now() - err = m.Migrate(MaxMigration) - log.Printf("[db] migration took %v\n", time.Since(startTime)) - if err != nil { - return err + for newVersion := curVersion + 1; newVersion <= targetVersion; newVersion++ { + err = MigrateUpStep(m, newVersion) + if err != nil { + return fmt.Errorf("during migration v%d: %w", err, newVersion) + } } + log.Printf("[db] migration done, new version = %d\n", targetVersion) return nil } -func MigrateVersion() (uint, bool, error) { - m, err := MakeMigrate() - if err != nil { - return 0, false, err +// returns curVersion, dirty, error +func MigrateVersion(m *migrate.Migrate) (uint, bool, error) { + if m == nil { + var err error + m, err = MakeMigrate() + if err != nil { + return 0, false, err + } } - return m.Version() + curVersion, dirty, err := m.Version() + if err == migrate.ErrNilVersion { + return 0, false, nil + } + return curVersion, dirty, err } func MigrateDown() error { @@ -111,6 +133,13 @@ func MigrateDown() error { } func MigrateGoto(n uint) error { + curVersion, _, _ := MigrateVersion(nil) + if curVersion == n { + return nil + } + if curVersion < n { + return MigrateUp(n) + } m, err := MakeMigrate() if err != nil { return err @@ -123,10 +152,12 @@ func MigrateGoto(n uint) error { } func TryMigrateUp() error { - err := MigrateUp() - if err != nil && err.Error() == migrate.ErrNoChange.Error() { - err = nil + curVersion, _, _ := MigrateVersion(nil) + log.Printf("[db] db version = %d\n", curVersion) + if curVersion >= MaxMigration { + return nil } + err := MigrateUp(MaxMigration) if err != nil { return err } @@ -134,7 +165,7 @@ func TryMigrateUp() error { } func MigratePrintVersion() error { - version, dirty, err := MigrateVersion() + version, dirty, err := MigrateVersion(nil) if err != nil { return fmt.Errorf("error getting db version: %v", err) } @@ -150,7 +181,7 @@ func MigrateCommandOpts(opts []string) error { if opts[0] == "--migrate-up" { fmt.Printf("migrate-up %v\n", GetDBName()) time.Sleep(3 * time.Second) - err = MigrateUp() + err = MigrateUp(MaxMigration) } else if opts[0] == "--migrate-down" { fmt.Printf("migrate-down %v\n", GetDBName()) time.Sleep(3 * time.Second) diff --git a/pkg/sstore/sstore.go b/pkg/sstore/sstore.go index 433f31ad..c10a210a 100644 --- a/pkg/sstore/sstore.go +++ b/pkg/sstore/sstore.go @@ -17,13 +17,13 @@ import ( "sync" "time" - "github.com/google/uuid" - "github.com/jmoiron/sqlx" - "github.com/sawka/txwrap" "github.com/commandlinedev/apishell/pkg/packet" "github.com/commandlinedev/apishell/pkg/shexec" "github.com/commandlinedev/prompt-server/pkg/dbutil" "github.com/commandlinedev/prompt-server/pkg/scbase" + "github.com/google/uuid" + "github.com/jmoiron/sqlx" + "github.com/sawka/txwrap" _ "github.com/mattn/go-sqlite3" ) @@ -239,28 +239,20 @@ type FeOptsType struct { TermFontSize int `json:"termfontsize,omitempty"` } -type ClientMigrationData struct { - MigrationType string `json:"migrationtype"` - MigrationPos int `json:"migrationpos"` - MigrationTotal int `json:"migrationtotal"` - MigrationDone bool `json:"migrationdone"` -} - type ClientData struct { - ClientId string `json:"clientid"` - UserId string `json:"userid"` - UserPrivateKeyBytes []byte `json:"-"` - UserPublicKeyBytes []byte `json:"-"` - UserPrivateKey *ecdsa.PrivateKey `json:"-" dbmap:"-"` - UserPublicKey *ecdsa.PublicKey `json:"-" dbmap:"-"` - ActiveSessionId string `json:"activesessionid"` - WinSize ClientWinSizeType `json:"winsize"` - ClientOpts ClientOptsType `json:"clientopts"` - FeOpts FeOptsType `json:"feopts"` - CmdStoreType string `json:"cmdstoretype"` - Migration *ClientMigrationData `json:"migration,omitempty" dbmap:"-"` - DBVersion int `json:"dbversion" dbmap:"-"` - OpenAIOpts *OpenAIOptsType `json:"openaiopts,omitempty" dbmap:"openaiopts"` + ClientId string `json:"clientid"` + UserId string `json:"userid"` + UserPrivateKeyBytes []byte `json:"-"` + UserPublicKeyBytes []byte `json:"-"` + UserPrivateKey *ecdsa.PrivateKey `json:"-" dbmap:"-"` + UserPublicKey *ecdsa.PublicKey `json:"-" dbmap:"-"` + ActiveSessionId string `json:"activesessionid"` + WinSize ClientWinSizeType `json:"winsize"` + ClientOpts ClientOptsType `json:"clientopts"` + FeOpts FeOptsType `json:"feopts"` + CmdStoreType string `json:"cmdstoretype"` + DBVersion int `json:"dbversion" dbmap:"-"` + OpenAIOpts *OpenAIOptsType `json:"openaiopts,omitempty" dbmap:"openaiopts"` } func (ClientData) UseDBMap() {} @@ -1220,8 +1212,8 @@ func createClientData(tx *TxWrap) error { WinSize: ClientWinSizeType{}, CmdStoreType: CmdStoreTypeScreen, } - query := `INSERT INTO client ( clientid, userid, activesessionid, userpublickeybytes, userprivatekeybytes, winsize) - VALUES (:clientid,:userid,:activesessionid,:userpublickeybytes,:userprivatekeybytes,:winsize)` + query := `INSERT INTO client ( clientid, userid, activesessionid, userpublickeybytes, userprivatekeybytes, winsize, cmdstoretype) + VALUES (:clientid,:userid,:activesessionid,:userpublickeybytes,:userprivatekeybytes,:winsize,:cmdstoretype)` tx.NamedExec(query, dbutil.ToDBMap(c, false)) log.Printf("create new clientid[%s] userid[%s] with public/private keypair\n", c.ClientId, c.UserId) return nil @@ -1273,28 +1265,6 @@ func EnsureClientData(ctx context.Context) (*ClientData, error) { return rtn, nil } -func GetCmdMigrationInfo(ctx context.Context) (*ClientMigrationData, error) { - return WithTxRtn(ctx, func(tx *TxWrap) (*ClientMigrationData, error) { - cdata := dbutil.GetMappable[*ClientData](tx, `SELECT * FROM client`) - if cdata == nil { - return nil, fmt.Errorf("no client data found") - } - if cdata.CmdStoreType == "session" { - total := tx.GetInt(`SELECT count(*) FROM cmd`) - posInv := tx.GetInt(`SELECT count(*) FROM cmd_migrate`) - mdata := &ClientMigrationData{ - MigrationType: "cmdscreen", - MigrationPos: total - posInv, - MigrationTotal: total, - MigrationDone: false, - } - return mdata, nil - } - // no migration info - return nil, nil - }) -} - func SetClientOpts(ctx context.Context, clientOpts ClientOptsType) error { txErr := WithTx(ctx, func(tx *TxWrap) error { query := `UPDATE client SET clientopts = ?` @@ -1303,93 +1273,3 @@ func SetClientOpts(ctx context.Context, clientOpts ClientOptsType) error { }) return txErr } - -type cmdMigrationType struct { - SessionId string - ScreenId string - CmdId string -} - -func getSliceChunk[T any](slice []T, chunkSize int) ([]T, []T) { - if chunkSize >= len(slice) { - return slice, nil - } - return slice[0:chunkSize], slice[chunkSize:] -} - -func processChunk(ctx context.Context, mchunk []cmdMigrationType) error { - for _, mig := range mchunk { - newFile, err := scbase.PtyOutFile(mig.ScreenId, mig.CmdId) - if err != nil { - log.Printf("ptyoutfile error: %v\n", err) - continue - } - oldFile, err := scbase.PtyOutFile_Sessions(mig.SessionId, mig.CmdId) - if err != nil { - log.Printf("ptyoutfile_sessions error: %v\n", err) - continue - } - err = os.Rename(oldFile, newFile) - if err != nil { - log.Printf("error renaming %s => %s: %v\n", oldFile, newFile, err) - continue - } - } - txErr := WithTx(ctx, func(tx *TxWrap) error { - for _, mig := range mchunk { - query := `DELETE FROM cmd_migrate WHERE cmdid = ?` - tx.Exec(query, mig.CmdId) - } - return nil - }) - if txErr != nil { - return txErr - } - return nil -} - -func RunCmdScreenMigration() { - ctx := context.Background() - startTime := time.Now() - mdata, err := GetCmdMigrationInfo(ctx) - if err != nil { - log.Printf("[prompt] error trying to run cmd migration: %v\n", err) - return - } - if mdata == nil || mdata.MigrationType != "cmdscreen" { - return - } - var migrations []cmdMigrationType - txErr := WithTx(ctx, func(tx *TxWrap) error { - tx.Select(&migrations, `SELECT * FROM cmd_migrate`) - return nil - }) - if txErr != nil { - log.Printf("[prompt] error trying to get cmd migrations: %v\n", txErr) - return - } - log.Printf("[db] got %d cmd migrations\n", len(migrations)) - for len(migrations) > 0 { - var mchunk []cmdMigrationType - mchunk, migrations = getSliceChunk(migrations, 5) - err = processChunk(ctx, mchunk) - if err != nil { - log.Printf("[prompt] cmd migration failed on chunk: %v\n%#v\n", err, mchunk) - return - } - } - err = os.RemoveAll(scbase.GetSessionsDir()) - if err != nil { - log.Printf("[db] cannot remove old sessions dir %s: %v\n", scbase.GetSessionsDir(), err) - } - txErr = WithTx(ctx, func(tx *TxWrap) error { - query := `UPDATE client SET cmdstoretype = 'screen'` - tx.Exec(query) - return nil - }) - if txErr != nil { - log.Printf("[db] cannot change client cmdstoretype: %v\n", err) - } - log.Printf("[db] cmd screen migration done: %v\n", time.Since(startTime)) - return -} diff --git a/pkg/sstore/sstore_migrate.go b/pkg/sstore/sstore_migrate.go new file mode 100644 index 00000000..909ff464 --- /dev/null +++ b/pkg/sstore/sstore_migrate.go @@ -0,0 +1,91 @@ +package sstore + +import ( + "context" + "fmt" + "log" + "os" + "time" + + "github.com/commandlinedev/prompt-server/pkg/scbase" +) + +type cmdMigrationType struct { + SessionId string + ScreenId string + CmdId string +} + +func getSliceChunk[T any](slice []T, chunkSize int) ([]T, []T) { + if chunkSize >= len(slice) { + return slice, nil + } + return slice[0:chunkSize], slice[chunkSize:] +} + +func RunCmdScreenMigration13() error { + ctx := context.Background() + startTime := time.Now() + var migrations []cmdMigrationType + txErr := WithTx(ctx, func(tx *TxWrap) error { + tx.Select(&migrations, `SELECT * FROM cmd_migrate`) + return nil + }) + if txErr != nil { + return fmt.Errorf("trying to get cmd migrations: %w", txErr) + } + log.Printf("[db] got %d cmd-screen migrations\n", len(migrations)) + for len(migrations) > 0 { + var mchunk []cmdMigrationType + mchunk, migrations = getSliceChunk(migrations, 5) + err := processMigrationChunk(ctx, mchunk) + if err != nil { + return fmt.Errorf("cmd migration failed on chunk: %w", err) + } + } + err := os.RemoveAll(scbase.GetSessionsDir()) + if err != nil { + return fmt.Errorf("cannot remove old sessions dir %s: %w\n", scbase.GetSessionsDir(), err) + } + txErr = WithTx(ctx, func(tx *TxWrap) error { + query := `UPDATE client SET cmdstoretype = 'screen'` + tx.Exec(query) + return nil + }) + if txErr != nil { + return fmt.Errorf("cannot change client cmdstoretype: %w", err) + } + log.Printf("[db] cmd screen migration done: %v\n", time.Since(startTime)) + return nil +} + +func processMigrationChunk(ctx context.Context, mchunk []cmdMigrationType) error { + for _, mig := range mchunk { + newFile, err := scbase.PtyOutFile(mig.ScreenId, mig.CmdId) + if err != nil { + log.Printf("ptyoutfile error: %v\n", err) + continue + } + oldFile, err := scbase.PtyOutFile_Sessions(mig.SessionId, mig.CmdId) + if err != nil { + log.Printf("ptyoutfile_sessions error: %v\n", err) + continue + } + err = os.Rename(oldFile, newFile) + if err != nil { + log.Printf("error renaming %s => %s: %v\n", oldFile, newFile, err) + continue + } + } + txErr := WithTx(ctx, func(tx *TxWrap) error { + for _, mig := range mchunk { + query := `DELETE FROM cmd_migrate WHERE cmdid = ?` + tx.Exec(query, mig.CmdId) + } + return nil + }) + if txErr != nil { + return txErr + } + return nil +}