mirror of
https://github.com/usememos/memos.git
synced 2025-06-05 22:09:59 +02:00
chore: fix postgres stmts
This commit is contained in:
Steven
@ -2,132 +2,113 @@ package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Masterminds/squirrel"
|
||||
"github.com/pkg/errors"
|
||||
"strings"
|
||||
|
||||
"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}
|
||||
|
||||
builder = builder.Values(values...)
|
||||
builder = builder.Suffix("RETURNING id")
|
||||
|
||||
// Prepare the final query
|
||||
query, args, err := builder.ToSql()
|
||||
if err != nil {
|
||||
fields := []string{"username", "role", "email", "nickname", "password_hash", "avatar_url"}
|
||||
args := []any{create.Username, create.Role, create.Email, create.Nickname, create.PasswordHash, create.AvatarURL}
|
||||
stmt := "INSERT INTO \"user\" (" + strings.Join(fields, ", ") + ") VALUES (" + placeholders(len(args)) + ") RETURNING id, avatar_url, created_ts, updated_ts, row_status"
|
||||
if err := d.db.QueryRowContext(ctx, stmt, args...).Scan(
|
||||
&create.ID,
|
||||
&create.AvatarURL,
|
||||
&create.CreatedTs,
|
||||
&create.UpdatedTs,
|
||||
&create.RowStatus,
|
||||
); 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
|
||||
return create, 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
|
||||
set, args := []string{}, []any{}
|
||||
if v := update.UpdatedTs; v != nil {
|
||||
builder = builder.Set("updated_ts", *v)
|
||||
set, args = append(set, "updated_ts = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := update.RowStatus; v != nil {
|
||||
builder = builder.Set("row_status", *v)
|
||||
set, args = append(set, "row_status = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := update.Username; v != nil {
|
||||
builder = builder.Set("username", *v)
|
||||
set, args = append(set, "username = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := update.Email; v != nil {
|
||||
builder = builder.Set("email", *v)
|
||||
set, args = append(set, "email = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := update.Nickname; v != nil {
|
||||
builder = builder.Set("nickname", *v)
|
||||
set, args = append(set, "nickname = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := update.AvatarURL; v != nil {
|
||||
builder = builder.Set("avatar_url", *v)
|
||||
set, args = append(set, "avatar_url = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := update.PasswordHash; v != nil {
|
||||
builder = builder.Set("password_hash", *v)
|
||||
set, args = append(set, "password_hash = "+placeholder(len(args)+1)), append(args, *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 {
|
||||
query := `
|
||||
UPDATE "user"
|
||||
SET ` + strings.Join(set, ", ") + `
|
||||
WHERE id = ` + placeholder(len(args)+1) + `
|
||||
RETURNING id, username, role, email, nickname, password_hash, avatar_url, created_ts, updated_ts, row_status
|
||||
`
|
||||
args = append(args, update.ID)
|
||||
user := &store.User{}
|
||||
if err := d.db.QueryRowContext(ctx, query, args...).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
|
||||
}
|
||||
|
||||
// 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", "created_ts", "updated_ts", "row_status").
|
||||
From(`"user"`).
|
||||
PlaceholderFormat(squirrel.Dollar)
|
||||
where, args := []string{"1 = 1"}, []any{}
|
||||
|
||||
// 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})
|
||||
where, args = append(where, "id = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := find.Username; v != nil {
|
||||
builder = builder.Where(squirrel.Eq{"username": *v})
|
||||
where, args = append(where, "username = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := find.Role; v != nil {
|
||||
builder = builder.Where(squirrel.Eq{"role": *v})
|
||||
where, args = append(where, "role = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := find.Email; v != nil {
|
||||
builder = builder.Where(squirrel.Eq{"email": *v})
|
||||
where, args = append(where, "email = "+placeholder(len(args)+1)), append(args, *v)
|
||||
}
|
||||
if v := find.Nickname; v != nil {
|
||||
builder = builder.Where(squirrel.Eq{"nickname": *v})
|
||||
where, args = append(where, "nickname = "+placeholder(len(args)+1)), append(args, *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
|
||||
query := `
|
||||
SELECT
|
||||
id,
|
||||
username,
|
||||
role,
|
||||
email,
|
||||
nickname,
|
||||
password_hash,
|
||||
avatar_url,
|
||||
created_ts,
|
||||
updated_ts,
|
||||
row_status
|
||||
FROM "user"
|
||||
WHERE ` + strings.Join(where, " AND ") + `
|
||||
ORDER BY created_ts DESC, row_status DESC
|
||||
`
|
||||
rows, err := d.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@ -161,35 +142,13 @@ func (d *DB) ListUsers(ctx context.Context, find *store.FindUser) ([]*store.User
|
||||
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()
|
||||
result, err := d.db.ExecContext(ctx, `
|
||||
DELETE FROM "user" WHERE id = $1
|
||||
`, delete.ID)
|
||||
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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user