go-socket/database.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 clients (
		    id SERIAL PRIMARY KEY,
		    name VARCHAR(20) UNIQUE NOT NULL,
		    pass_hash VARCHAR(60) NOT NULL,
		    pronouns VARCHAR(15) 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 chat_groups (
		    id SERIAL PRIMARY KEY,
		    name VARCHAR(48) NOT NULL,
		    creator_id INTEGER NOT NULL REFERENCES clients(id) ON DELETE CASCADE,
		    owner_id INTEGER NOT NULL REFERENCES clients(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 clients(id) ON DELETE CASCADE,
			joined_at TIMESTAMP NOT NULL DEFAULT NOW(),
			PRIMARY KEY (group_id, user_id)
		)
	`)
	if err != nil {
		panic(err)
	}
}

func DbSaveClientWithoutGroups(ctx context.Context, client *Client) error {
	err := dbConn.QueryRow(ctx, `
        INSERT INTO clients (name, pass_hash, pronouns, color_red, color_green, color_blue, created_at)
        VALUES ($1, $2, $3, $4, $5, $6, $7)
        RETURNING id
    `, client.Name, client.PasswordHash, client.Pronouns, client.Color[0], client.Color[1], client.Color[2], client.CreatedAt).
		Scan(&client.Id)
	return err
}

func DbSetClientByName(ctx context.Context, client *Client) error {
	err := dbConn.QueryRow(ctx, `
		SELECT id, name, pass_hash, pronouns, color_red, color_green, color_blue, created_at FROM clients WHERE name = $1
	`, client.Name).Scan(&client.Id, &client.Name, &client.PasswordHash, &client.Pronouns, &client.Color[0], &client.Color[1], &client.Color[2], &client.CreatedAt)
	if err != nil {
		return err
	}
	return DbSetClientGroups(ctx, client)
}

func DbSetClientByIdWithoutGroups(ctx context.Context, client *Client) error {
	err := dbConn.QueryRow(ctx, `
		SELECT name, pass_hash, pronouns, color_red, color_green, color_blue, created_at FROM clients WHERE id = $1
	`, client.Id).Scan(&client.Name, &client.PasswordHash, &client.Pronouns, &client.Color[0], &client.Color[1], &client.Color[2], &client.CreatedAt)
	return err
}

func DbSetClientById(ctx context.Context, client *Client) error {
	err := DbSetClientByIdWithoutGroups(ctx, client)
	if err != nil {
		return err
	}
	return DbSetClientGroups(ctx, client)
}

func DbSaveGroupWithoutClients(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.EnableClientColors, 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
}

func DbSetGroupByIdWithoutClients(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.EnableClientColors, &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.Clients = make(map[uint32]struct{})
	for rows.Next() {
		var userId uint32
		if err := rows.Scan(&userId); err != nil {
			return err
		}
		group.Clients[userId] = struct{}{}
	}
	return rows.Err()
}

func DbSetGroupById(ctx context.Context, group *Group) error {
	err := DbSetGroupByIdWithoutClients(ctx, group)
	if err != nil {
		return err
	}
	return DbSetGroupMemberClients(ctx, group)
}

func DbSetGroupMemberClients(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.Clients = make(map[uint32]struct{})
	for rows.Next() {
		var userId uint32
		if err := rows.Scan(&userId); err != nil {
			return err
		}
		group.Clients[userId] = struct{}{}
	}
	return rows.Err()
}

func DbAddClientsToGroup(ctx context.Context, groupId uint32, clientIds []uint32) error {
	batch := &pgx.Batch{}
	now := time.Now()
	for _, cid := range clientIds {
		batch.Queue(`
			INSERT INTO chat_group_members (group_id, user_id, joined_at)
			VALUES ($1, $2, $3)
			ON CONFLICT DO NOTHING
		`, groupId, cid, now)
	}
	br := dbConn.SendBatch(ctx, batch)
	defer br.Close()
	for range clientIds {
		if _, err := br.Exec(); err != nil {
			return err
		}
	}
	return nil
}

func DbSetClientGroups(ctx context.Context, client *Client) error {
	rows, err := dbConn.Query(ctx, `
		SELECT group_id FROM chat_group_members WHERE user_id = $1
	`, client.Id)
	if err != nil {
		return err
	}
	defer rows.Close()

	client.Groups = make(map[uint32]struct{})
	for rows.Next() {
		var groupId uint32
		if err := rows.Scan(&groupId); err != nil {
			return err
		}
		client.Groups[groupId] = struct{}{}
	}
	return rows.Err()
}