package psql import ( "database/sql" "errors" "github.com/vinovest/sqlx" ) type RPGeneralPreset struct { ID string Name string Description string PreHistory string `db:"pre_history"` PostHistory string `db:"post_history"` } type RPScenarios struct { ID int Name string Prompt string } type RPUser struct { UserID int64 `db:"user_id"` UserPrompt string `db:"user_prompt"` SelectedPreset string `db:"selected_preset"` UsedTokens int64 `db:"used_tokens"` } type RPRepository struct { db *sqlx.DB } func NewRPRepository(db *sqlx.DB) *RPRepository { return &RPRepository{db: db} } func (rep *RPRepository) GetOrCreateUser(id int64) (*RPUser, error) { user, err := rep.GetUser(id) if errors.Is(err, sql.ErrNoRows) { user, err = rep.CreateUser(id) } return user, err } func (rep *RPRepository) CreateUser(id int64) (*RPUser, error) { user := new(RPUser) err := rep.db.Get(user, "INSERT INTO rp_users(user_id) VALUES ($1) RETURNING *;", id) return user, err } func (rep *RPRepository) GetUser(id int64) (*RPUser, error) { user := new(RPUser) err := rep.db.Get(user, "SELECT * FROM rp_users WHERE user_id=$1", id) return user, err } func (rep *RPRepository) UpdateUser(user *RPUser) error { _, err := rep.db.NamedExec( "UPDATE rp_users SET selected_preset=:selected_preset, used_tokens=:used_tokens, user_prompt=:user_prompt WHERE user_id=:user_id;", user, ) return err } func (rep *RPRepository) UpdateUserPreset(user *RPUser, presetId string) (*RPGeneralPreset, error) { preset, err := rep.GetPreset(presetId) if err != nil { return preset, err } _, err = rep.db.Exec("UPDATE rp_users SET selected_preset=$1 WHERE user_id=$2;", presetId, user.UserID) return preset, err } func (rep *RPRepository) GetUserPreset(user *RPUser) (*RPGeneralPreset, error) { preset, err := rep.GetPreset(user.SelectedPreset) if errors.Is(err, sql.ErrNoRows) { return rep.UpdateUserPreset(user, "soft") } return preset, err } func (rep *RPRepository) GetAllPresets() ([]*RPGeneralPreset, error) { presets := make([]*RPGeneralPreset, 0) err := rep.db.Select(&presets, "SELECT * FROM rp_general_presets;") return presets, err } func (rep *RPRepository) GetPreset(id string) (*RPGeneralPreset, error) { preset := new(RPGeneralPreset) err := rep.db.Get(preset, "SELECT * FROM rp_general_presets WHERE id=$1;", id) return preset, err } func (rep *RPRepository) GetAllScenarios() ([]*RPScenarios, error) { scenarios := make([]*RPScenarios, 0) err := rep.db.Select(&scenarios, "SELECT * FROM rp_scenarios;") return scenarios, err } func (rep *RPRepository) GetScenario(id int) (*RPScenarios, error) { scenario := new(RPScenarios) err := rep.db.Get(scenario, "SELECT * FROM rp_scenarios WHERE id=$1;", id) return scenario, err }