read-only, write db instance more

Signed-off-by: Gyuho Lee <[email protected]>
leptonai · Jan 8, 2025 · ec38861 · ec38861
1 parent 99f090d
commit ec38861
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 24 deletions.
diff --git a/cmd/gpud/command/login.go b/cmd/gpud/command/login.go
@@ -31,26 +31,33 @@ func cmdLogin(cliContext *cli.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to get state file: %w", err)
 	}
-	db, err := sqlite.Open(stateFile)
+
+	dbRW, err := sqlite.Open(stateFile)
+	if err != nil {
+		return fmt.Errorf("failed to open state file: %w", err)
+	}
+	defer dbRW.Close()
+
+	dbRO, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
 	if err != nil {
 		return fmt.Errorf("failed to open state file: %w", err)
 	}
-	defer db.Close()
+	defer dbRO.Close()
 
-	uid, err := state.CreateMachineIDIfNotExist(rootCtx, db, "")
+	uid, err := state.CreateMachineIDIfNotExist(rootCtx, dbRW, dbRO, "")
 	if err != nil {
 		return fmt.Errorf("failed to get machine uid: %w", err)
 	}
 
-	components, err := state.GetComponents(rootCtx, db, uid)
+	components, err := state.GetComponents(rootCtx, dbRO, uid)
 	if err != nil {
 		return fmt.Errorf("failed to get components: %w", err)
 	}
 
 	cliToken := cliContext.String("token")
 	endpoint := cliContext.String("endpoint")
 
-	dbToken, _ := state.GetLoginInfo(rootCtx, db, uid)
+	dbToken, _ := state.GetLoginInfo(rootCtx, dbRO, uid)
 	token := dbToken
 	if cliToken != "" {
 		token = cliToken
@@ -88,7 +95,7 @@ func cmdLogin(cliContext *cli.Context) error {
 	}
 
 	if token != dbToken {
-		if err = state.UpdateLoginInfo(rootCtx, db, uid, token); err != nil {
+		if err = state.UpdateLoginInfo(rootCtx, dbRW, uid, token); err != nil {
 			fmt.Println("machine logged in but failed to update token:", err)
 		}
 	}

diff --git a/cmd/gpud/command/utils.go b/cmd/gpud/command/utils.go
@@ -14,12 +14,20 @@ func GetUID(ctx context.Context) (string, error) {
 	if err != nil {
 		return "", fmt.Errorf("failed to get state file: %w", err)
 	}
-	db, err := sqlite.Open(stateFile)
+
+	dbRW, err := sqlite.Open(stateFile)
+	if err != nil {
+		return "", fmt.Errorf("failed to open state file: %w", err)
+	}
+	defer dbRW.Close()
+
+	dbRO, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
 	if err != nil {
 		return "", fmt.Errorf("failed to open state file: %w", err)
 	}
-	defer db.Close()
-	uid, err := state.CreateMachineIDIfNotExist(ctx, db, "")
+	defer dbRO.Close()
+
+	uid, err := state.CreateMachineIDIfNotExist(ctx, dbRW, dbRO, "")
 	if err != nil {
 		return "", fmt.Errorf("failed to get machine uid: %w", err)
 	}

diff --git a/components/state/state.go b/components/state/state.go
@@ -25,8 +25,8 @@ const (
 	ColumnComponents  = "components"
 )
 
-func CreateTableMachineMetadata(ctx context.Context, db *sql.DB) error {
-	_, err := db.ExecContext(ctx, fmt.Sprintf(`
+func CreateTableMachineMetadata(ctx context.Context, dbRW *sql.DB) error {
+	_, err := dbRW.ExecContext(ctx, fmt.Sprintf(`
 CREATE TABLE IF NOT EXISTS %s (
 	%s TEXT PRIMARY KEY,
 	%s INTEGER,
@@ -36,7 +36,7 @@ CREATE TABLE IF NOT EXISTS %s (
 	return err
 }
 
-func CreateMachineIDIfNotExist(ctx context.Context, db *sql.DB, providedUID string) (string, error) {
+func CreateMachineIDIfNotExist(ctx context.Context, dbRW *sql.DB, dbRO *sql.DB, providedUID string) (string, error) {
 	query := fmt.Sprintf(`
 SELECT %s, %s FROM %s
 LIMIT 1;
@@ -50,7 +50,7 @@ LIMIT 1;
 		machineID   string
 		unixSeconds int64
 	)
-	err := db.QueryRowContext(ctx, query).Scan(&machineID, &unixSeconds)
+	err := dbRO.QueryRowContext(ctx, query).Scan(&machineID, &unixSeconds)
 	if err == nil { // reuse existing machine ID
 		return machineID, nil
 	}
@@ -81,13 +81,13 @@ INSERT OR REPLACE INTO %s (%s, %s) VALUES (?, ?);
 		ColumnMachineID,
 		ColumnUnixSeconds,
 	)
-	if _, err := db.ExecContext(ctx, query, uid, time.Now().UTC().Unix()); err != nil {
+	if _, err := dbRW.ExecContext(ctx, query, uid, time.Now().UTC().Unix()); err != nil {
 		return "", err
 	}
 	return uid, nil
 }
 
-func GetLoginInfo(ctx context.Context, db *sql.DB, machineID string) (string, error) {
+func GetLoginInfo(ctx context.Context, dbRO *sql.DB, machineID string) (string, error) {
 	query := fmt.Sprintf(`
 SELECT %s FROM %s WHERE %s = '%s'
 LIMIT 1;
@@ -98,11 +98,11 @@ LIMIT 1;
 		machineID,
 	)
 	var token string
-	err := db.QueryRowContext(ctx, query).Scan(&token)
+	err := dbRO.QueryRowContext(ctx, query).Scan(&token)
 	return token, err
 }
 
-func UpdateLoginInfo(ctx context.Context, db *sql.DB, machineID string, token string) error {
+func UpdateLoginInfo(ctx context.Context, dbRW *sql.DB, machineID string, token string) error {
 	query := fmt.Sprintf(`
 UPDATE %s SET %s = '%s' WHERE %s = '%s';
 `,
@@ -112,7 +112,7 @@ UPDATE %s SET %s = '%s' WHERE %s = '%s';
 		ColumnMachineID,
 		machineID,
 	)
-	if _, err := db.ExecContext(ctx, query); err != nil {
+	if _, err := dbRW.ExecContext(ctx, query); err != nil {
 		return err
 	}
 	return nil

diff --git a/components/state/state_test.go b/components/state/state_test.go
@@ -11,24 +11,24 @@ import (
 func TestOpenMemory(t *testing.T) {
 	t.Parallel()
 
-	db, err := sqlite.Open(":memory:")
+	dbRW, err := sqlite.Open(":memory:")
 	if err != nil {
 		t.Fatalf("failed to open database: %v", err)
 	}
-	defer db.Close()
+	defer dbRW.Close()
 
 	// create some table
-	if _, err = db.Exec("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)"); err != nil {
+	if _, err = dbRW.Exec("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)"); err != nil {
 		t.Fatal("failed to create table:", err)
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	if err := CreateTableMachineMetadata(ctx, db); err != nil {
+	if err := CreateTableMachineMetadata(ctx, dbRW); err != nil {
 		t.Fatal("failed to create table:", err)
 	}
-	id, err := CreateMachineIDIfNotExist(ctx, db, "")
+	id, err := CreateMachineIDIfNotExist(ctx, dbRW, dbRW, "")
 	if err != nil {
 		t.Fatal("failed to create machine id:", err)
 	}

diff --git a/internal/server/server.go b/internal/server/server.go
@@ -1215,7 +1215,7 @@ func New(ctx context.Context, config *lepconfig.Config, endpoint string, cliUID
 		}
 	}
 
-	uid, err := state.CreateMachineIDIfNotExist(ctx, dbRW, cliUID)
+	uid, err := state.CreateMachineIDIfNotExist(ctx, dbRW, dbRO, cliUID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create machine uid: %w", err)
 	}