Files
go-socket/database.go
T

378 lines
11 KiB
Go

package main
import (
"context"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
var dbConn *pgxpool.Pool
func DbInit(ctx context.Context) {
var err error
dbConn, err = pgxpool.New(ctx, "postgres://master:secret@localhost:5432") // TODO change to env in production
if err != nil {
panic(err)
}
_, err = dbConn.Exec(ctx, `
CREATE TABLE IF NOT EXISTS users (
id SERIAL PRIMARY KEY,
name TEXT UNIQUE NOT NULL,
pass_hash TEXT NOT NULL,
pronouns TEXT DEFAULT NULL,
color_red SMALLINT DEFAULT NULL,
color_green SMALLINT DEFAULT NULL,
color_blue SMALLINT DEFAULT NULL,
created_at TIMESTAMP NOT NULL DEFAULT NOW()
)
`)
if err != nil {
panic(err)
}
_, err = dbConn.Exec(ctx, `
CREATE TABLE IF NOT EXISTS user_connections (
requestor_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
recipient_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
is_accepted BOOLEAN NOT NULL DEFAULT FALSE,
created_at TIMESTAMP NOT NULL DEFAULT NOW()
)
`)
_, err = dbConn.Exec(ctx, `
CREATE TABLE IF NOT EXISTS chat_groups (
id SERIAL PRIMARY KEY,
name TEXT NOT NULL,
creator_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
owner_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
enable_client_colors BOOLEAN NOT NULL DEFAULT true,
color_red SMALLINT DEFAULT NULL,
color_green SMALLINT DEFAULT NULL,
color_blue SMALLINT DEFAULT NULL,
created_at TIMESTAMP NOT NULL DEFAULT NOW()
)
`)
if err != nil {
panic(err)
}
_, err = dbConn.Exec(ctx, `
CREATE TABLE IF NOT EXISTS chat_group_members (
group_id INTEGER NOT NULL REFERENCES chat_groups(id) ON DELETE CASCADE,
user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
joined_at TIMESTAMP NOT NULL DEFAULT NOW(),
PRIMARY KEY (group_id, user_id)
)
`)
if err != nil {
panic(err)
}
}
// DbUserSaveWithoutGroupsConnections saves user in db without groups and sets its id
// return: error if not successful
func DbUserSaveWithoutGroupsConnections(ctx context.Context, user *User) error {
err := dbConn.QueryRow(ctx, `
INSERT INTO users (name, pass_hash, pronouns, color_red, color_green, color_blue, created_at)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id
`, user.Name, user.PasswordHash, user.Pronouns, user.Color[0], user.Color[1], user.Color[2], user.CreatedAt).
Scan(&user.Id)
return err
}
// DbUserDelete deletes given user by id
// return: error if not successful
func DbUserDelete(ctx context.Context, id uint32) error {
_, err := dbConn.Exec(ctx, `
DELETE FROM users WHERE id = $1
`, id)
return err
}
// DbUserSetByName sets all fields of given struct with database's data
// return: error if not successful
func DbUserSetByName(ctx context.Context, user *User) error {
err := dbConn.QueryRow(ctx, `
SELECT id, name, pass_hash, pronouns, color_red, color_green, color_blue, created_at FROM users WHERE name = $1
`, user.Name).Scan(&user.Id, &user.Name, &user.PasswordHash, &user.Pronouns, &user.Color[0], &user.Color[1], &user.Color[2], &user.CreatedAt)
if err != nil {
return err
}
err = DbUserSetGroups(ctx, user)
if err != nil {
return err
}
return DbUserSetConnections(ctx, user)
}
// DbUserSetByIdWithoutGroupsConnections sets all fields of given struct with database's data using id, excluding groups
// return: error if not successful
func DbUserSetByIdWithoutGroupsConnections(ctx context.Context, user *User) error {
err := dbConn.QueryRow(ctx, `
SELECT name, pass_hash, pronouns, color_red, color_green, color_blue, created_at FROM users WHERE id = $1
`, user.Id).Scan(&user.Name, &user.PasswordHash, &user.Pronouns, &user.Color[0], &user.Color[1], &user.Color[2], &user.CreatedAt)
return err
}
// DbUserSetById sets all fields of given struct with database's data using id, including groups
// return: error if not successful
func DbUserSetById(ctx context.Context, user *User) error {
err := DbUserSetByIdWithoutGroupsConnections(ctx, user)
if err != nil {
return err
}
err = DbUserSetGroups(ctx, user)
if err != nil {
return err
}
return DbUserSetConnections(ctx, user)
}
// DbUserSetGroups populates user's Groups map with ids of all groups the user belongs to from database
// return: error if not successful
func DbUserSetGroups(ctx context.Context, user *User) error {
rows, err := dbConn.Query(ctx, `
SELECT group_id FROM chat_group_members WHERE user_id = $1
`, user.Id)
if err != nil {
return err
}
defer rows.Close()
user.Groups = make(map[uint32]struct{})
for rows.Next() {
var groupId uint32
if err := rows.Scan(&groupId); err != nil {
return err
}
user.Groups[groupId] = struct{}{}
}
return rows.Err()
}
func DbUserSetConnections(ctx context.Context, user *User) error {
rows, err := dbConn.Query(ctx, `
SELECT requestor_id, recipient_id, is_accepted, created_at FROM user_connections WHERE requestor_id = $1 or recipient_id = $1
`)
if err != nil {
return err
}
defer rows.Close()
for rows.Next() {
var (
requestorId uint32
recipientId uint32
isAccepted bool
createdAt time.Time
)
err = rows.Scan(&requestorId, &recipientId, &isAccepted, &createdAt)
if err != nil {
return err
}
if requestorId == user.Id {
user.Connections[recipientId] = &Connection{
CreatedAt: createdAt,
With: recipientId,
IsFromUser: true,
IsAccepted: isAccepted,
}
} else {
user.Connections[requestorId] = &Connection{
CreatedAt: createdAt,
With: requestorId,
IsFromUser: false,
IsAccepted: isAccepted,
}
}
}
return rows.Err()
}
// DbUserSetColor set user's color based on id
// return: error if not successful
func DbUserSetColor(ctx context.Context, user *User) error {
_, err := dbConn.Exec(ctx, `
UPDATE users SET color_red = $1, color_green = $2, color_blue = $3 WHERE id = $4
`, user.Color[0], user.Color[1], user.Color[2], user.Id)
return err
}
// DbUserSetPronouns set user's pronouns based on id
// return: error if not successful
func DbUserSetPronouns(ctx context.Context, user *User) error {
_, err := dbConn.Exec(ctx, `
UPDATE users SET pronouns = $1 WHERE id = $2
`, user.Pronouns, user.Id)
return err
}
// DbGroupSetColor set group's color based on id
// return: error if not successful
func DbGroupSetColor(ctx context.Context, group *Group) error {
_, err := dbConn.Exec(ctx, `
UPDATE chat_groups SET color_red = $1, color_green = $2, color_blue = $3 WHERE id = $4
`, group.Color[0], group.Color[1], group.Color[2], group.Id)
return err
}
// DbGroupSaveWithoutUsers saves group in db and sets its id, also adds the owner as first member
// return: error if not successful
func DbGroupSaveWithoutUsers(ctx context.Context, group *Group) error {
err := dbConn.QueryRow(ctx, `
INSERT INTO chat_groups (name, creator_id, owner_id, enable_client_colors, color_red, color_green, color_blue, created_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
RETURNING id
`, group.Name, group.CreatorId, group.OwnerId, group.EnableUserColors, group.Color[0], group.Color[1], group.Color[2], group.CreatedAt).
Scan(&group.Id)
if err != nil {
return err
}
_, err = dbConn.Exec(ctx, `
INSERT INTO chat_group_members (group_id, user_id, joined_at)
VALUES ($1, $2, $3)
`, group.Id, group.OwnerId, group.CreatedAt)
return err
}
// DbGroupDelete deletes given group by id
// return: error if not successful
func DbGroupDelete(ctx context.Context, group *Group) error {
_, err := dbConn.Exec(ctx, `
DELETE FROM chat_groups WHERE id = $1
`, group.Id)
return err
}
// DbGroupSetWithoutUsersById sets all fields of given struct with database's data using id, populates Users map with member ids but not their data
// return: error if not successful
func DbGroupSetWithoutUsersById(ctx context.Context, group *Group) error {
err := dbConn.QueryRow(ctx, `
SELECT name, creator_id, owner_id, enable_client_colors, color_red, color_green, color_blue, created_at FROM chat_groups WHERE id = $1
`, group.Id).Scan(&group.Name, &group.CreatorId, &group.OwnerId, &group.EnableUserColors, &group.Color[0], &group.Color[1], &group.Color[2], &group.CreatedAt)
if err != nil {
return err
}
rows, err := dbConn.Query(ctx, `
SELECT user_id FROM chat_group_members WHERE group_id = $1
`, group.Id)
if err != nil {
return err
}
defer rows.Close()
group.Users = make(map[uint32]struct{})
for rows.Next() {
var userId uint32
if err := rows.Scan(&userId); err != nil {
return err
}
group.Users[userId] = struct{}{}
}
return rows.Err()
}
// DbGroupSetById sets all fields of given struct with database's data using id, including full member user data
// return: error if not successful
func DbGroupSetById(ctx context.Context, group *Group) error {
err := DbGroupSetWithoutUsersById(ctx, group)
if err != nil {
return err
}
return DbGroupSetMembers(ctx, group)
}
// DbGroupSetMembers populates group's Users map with ids of all members from database
// return: error if not successful
func DbGroupSetMembers(ctx context.Context, group *Group) error {
rows, err := dbConn.Query(ctx, `
SELECT user_id FROM chat_group_members WHERE group_id = $1
`, group.Id)
if err != nil {
return err
}
defer rows.Close()
group.Users = make(map[uint32]struct{})
for rows.Next() {
var userId uint32
if err := rows.Scan(&userId); err != nil {
return err
}
group.Users[userId] = struct{}{}
}
return rows.Err()
}
// DbGroupAddUsers adds given users to group in db, silently ignores already existing members
// return: error if not successful
func DbGroupAddUsers(ctx context.Context, groupId uint32, userIds *[MaxUsersInGroup]uint32) error {
batch := &pgx.Batch{}
now := time.Now()
var count int
for _, uid := range userIds {
if uid == 0 {
continue
}
batch.Queue(`
INSERT INTO chat_group_members (group_id, user_id, joined_at)
VALUES ($1, $2, $3)
ON CONFLICT DO NOTHING
`, groupId, uid, now)
count++
}
br := dbConn.SendBatch(ctx, batch)
defer br.Close()
for range count {
if _, err := br.Exec(); err != nil {
return err
}
}
return nil
}
// DbGroupRemoveUsers removes given users from group in db, silently ignores not existing members
// return: deleted users count, error if not successful
func DbGroupRemoveUsers(ctx context.Context, groupId uint32, userIds *[MaxUsersInGroup]uint32) (int, error) {
batch := &pgx.Batch{}
var count int
for _, uid := range userIds {
if uid == 0 {
continue
}
batch.Queue(`
DELETE FROM chat_group_members WHERE group_id = $1 AND user_id = $2
`, groupId, uid)
count++
}
br := dbConn.SendBatch(ctx, batch)
defer br.Close()
var deleted int
for range count {
tag, err := br.Exec()
if err != nil {
return deleted, err
}
deleted += int(tag.RowsAffected())
}
return deleted, nil
}
// DbGroupSetOwnerId set group's owner based on id
// return: error if not successful
func DbGroupSetOwnerId(ctx context.Context, group *Group) error {
_, err := dbConn.Exec(ctx, `
UPDATE chat_groups SET owner_id = $1 WHERE id = $2
`, group.OwnerId, group.Id)
return err
}