package postgres import ( "context" "github.com/Masterminds/squirrel" "github.com/pkg/errors" "github.com/usememos/memos/store" ) func (d *DB) CreateUser(ctx context.Context, create *store.User) (*store.User, error) { // Start building the insert statement builder := squirrel.Insert("\"user\"").PlaceholderFormat(squirrel.Dollar) columns := []string{"username", "role", "email", "nickname", "password_hash", "avatar_url"} builder = builder.Columns(columns...) values := []any{create.Username, create.Role, create.Email, create.Nickname, create.PasswordHash, create.AvatarURL} if create.RowStatus != "" { builder = builder.Columns("row_status") values = append(values, create.RowStatus) } if create.CreatedTs != 0 { builder = builder.Columns("created_ts") values = append(values, squirrel.Expr("TO_TIMESTAMP(?)", create.CreatedTs)) } if create.UpdatedTs != 0 { builder = builder.Columns("updated_ts") values = append(values, squirrel.Expr("TO_TIMESTAMP(?)", create.UpdatedTs)) } if create.ID != 0 { builder = builder.Columns("id") values = append(values, create.ID) } builder = builder.Values(values...) builder = builder.Suffix("RETURNING id") // Prepare the final query query, args, err := builder.ToSql() if err != nil { return nil, err } // Execute the query and get the returned ID var id int32 err = d.db.QueryRowContext(ctx, query, args...).Scan(&id) if err != nil { return nil, err } // Use the returned ID to retrieve the full user object user, err := d.GetUser(ctx, &store.FindUser{ID: &id}) if err != nil { return nil, err } return user, nil } func (d *DB) UpdateUser(ctx context.Context, update *store.UpdateUser) (*store.User, error) { // Start building the update statement builder := squirrel.Update("\"user\"").PlaceholderFormat(squirrel.Dollar) // Conditionally add set clauses if v := update.UpdatedTs; v != nil { builder = builder.Set("updated_ts", squirrel.Expr("to_timestamp(?)", *v)) } if v := update.RowStatus; v != nil { builder = builder.Set("row_status", *v) } if v := update.Username; v != nil { builder = builder.Set("username", *v) } if v := update.Email; v != nil { builder = builder.Set("email", *v) } if v := update.Nickname; v != nil { builder = builder.Set("nickname", *v) } if v := update.AvatarURL; v != nil { builder = builder.Set("avatar_url", *v) } if v := update.PasswordHash; v != nil { builder = builder.Set("password_hash", *v) } // Add the WHERE clause builder = builder.Where(squirrel.Eq{"id": update.ID}) // Prepare the final query query, args, err := builder.ToSql() if err != nil { return nil, err } // Execute the query with the context if _, err := d.db.ExecContext(ctx, query, args...); err != nil { return nil, err } // Retrieve the updated user user, err := d.GetUser(ctx, &store.FindUser{ID: &update.ID}) if err != nil { return nil, err } return user, nil } func (d *DB) ListUsers(ctx context.Context, find *store.FindUser) ([]*store.User, error) { // Start building the SELECT statement builder := squirrel.Select("id", "username", "role", "email", "nickname", "password_hash", "avatar_url", "FLOOR(EXTRACT(EPOCH FROM created_ts)) AS created_ts", "FLOOR(EXTRACT(EPOCH FROM updated_ts)) AS updated_ts", "row_status"). From("\"user\""). PlaceholderFormat(squirrel.Dollar) // 1 = 1 is often used as a no-op in SQL, ensuring there's always a WHERE clause builder = builder.Where("1 = 1") // Conditionally add where clauses if v := find.ID; v != nil { builder = builder.Where(squirrel.Eq{"id": *v}) } if v := find.Username; v != nil { builder = builder.Where(squirrel.Eq{"username": *v}) } if v := find.Role; v != nil { builder = builder.Where(squirrel.Eq{"role": *v}) } if v := find.Email; v != nil { builder = builder.Where(squirrel.Eq{"email": *v}) } if v := find.Nickname; v != nil { builder = builder.Where(squirrel.Eq{"nickname": *v}) } // Add ordering builder = builder.OrderBy("created_ts DESC", "row_status DESC") // Prepare the final query query, args, err := builder.ToSql() if err != nil { return nil, err } // Execute the query with the context rows, err := d.db.QueryContext(ctx, query, args...) if err != nil { return nil, err } defer rows.Close() list := make([]*store.User, 0) for rows.Next() { var user store.User if err := rows.Scan( &user.ID, &user.Username, &user.Role, &user.Email, &user.Nickname, &user.PasswordHash, &user.AvatarURL, &user.CreatedTs, &user.UpdatedTs, &user.RowStatus, ); err != nil { return nil, err } list = append(list, &user) } if err := rows.Err(); err != nil { return nil, err } return list, nil } func (d *DB) GetUser(ctx context.Context, find *store.FindUser) (*store.User, error) { list, err := d.ListUsers(ctx, find) if err != nil { return nil, err } if len(list) != 1 { return nil, errors.Wrapf(nil, "unexpected user count: %d", len(list)) } return list[0], nil } func (d *DB) DeleteUser(ctx context.Context, delete *store.DeleteUser) error { // Start building the DELETE statement builder := squirrel.Delete("\"user\""). PlaceholderFormat(squirrel.Dollar). Where(squirrel.Eq{"id": delete.ID}) // Prepare the final query query, args, err := builder.ToSql() if err != nil { return err } // Execute the query with the context result, err := d.db.ExecContext(ctx, query, args...) if err != nil { return err } if _, err := result.RowsAffected(); err != nil { return err } if err := d.Vacuum(ctx); err != nil { // Prevent linter warning. return err } return nil }