go-socket/packages/postgresql/postgresql.go

package postgresql

import (
	"context"
	"fmt"
	"strings"
	"time"

	"go-socket/packages/cache"
	"go-socket/packages/convertions"
	"go-socket/packages/types"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

var dbConn *pgxpool.Pool

func Init(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 EXTENSION IF NOT EXISTS "pgcrypto";`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS users (
		    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
		    name TEXT UNIQUE NOT NULL,
		    pass_hash TEXT NOT NULL,
		    pronouns TEXT DEFAULT NULL,
			description TEXT DEFAULT NULL,
			avatar TEXT DEFAULT NULL,
			profile_bg TEXT DEFAULT NULL,
		    rgba BIGINT NOT NULL DEFAULT 0 CHECK (rgba BETWEEN 0 AND 4294967295),
			created_at TIMESTAMP NOT NULL DEFAULT NOW()
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS user_connections (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			requestor_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
			recipient_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
		    state SMALLINT NOT NULL DEFAULT 0,
			created_at TIMESTAMP NOT NULL DEFAULT NOW()
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS direct_messages (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			sender_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
			receiver_id UUID NOT NULL REFERENCES user_connections(id) ON DELETE CASCADE,
			created_at TIMESTAMP NOT NULL DEFAULT NOW(),
			content TEXT NOT NULL,
			attached_file TEXT NOT NULL DEFAULT ''
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS hubs (
			id UUID PRIMARY KEY,
			creator_id UUID NOT NULL REFERENCES users(id),
			name TEXT NOT NULL,
			allow_user_color BOOLEAN NOT NULL DEFAULT TRUE,
			rgba BIGINT NOT NULL DEFAULT 0 CHECK (rgba BETWEEN 0 AND 4294967295),
			created_at TIMESTAMP NOT NULL DEFAULT NOW()
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS hub_roles (
			hub_id UUID NOT NULL REFERENCES hubs(id) ON DELETE CASCADE,
			role_id SMALLINT NOT NULL,
			name TEXT NOT NULL,
			permissions INTEGER NOT NULL DEFAULT 0,
			rgba BIGINT NOT NULL DEFAULT 0 CHECK (rgba BETWEEN 0 AND 4294967295),
			PRIMARY KEY (hub_id, role_id)
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS hub_channel_groups (
			channel_group_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			hub_id UUID NOT NULL REFERENCES hubs(id) ON DELETE CASCADE,
			name TEXT NOT NULL,
			rgba BIGINT NOT NULL DEFAULT 0 CHECK (rgba BETWEEN 0 AND 4294967295),
			position SMALLINT NOT NULL DEFAULT 0
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS hub_channels (
			channel_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			hub_id UUID NOT NULL REFERENCES hubs(id) ON DELETE CASCADE,
			name TEXT NOT NULL,
			parent_group_id UUID REFERENCES hub_channel_groups(channel_group_id) ON DELETE CASCADE,
			position SMALLINT NOT NULL DEFAULT 0
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS hub_users (
			hub_id UUID NOT NULL REFERENCES hubs(id) ON DELETE CASCADE,
			user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
			name TEXT NOT NULL,
			created_at TIMESTAMP NOT NULL DEFAULT NOW(),
			PRIMARY KEY (hub_id, user_id)
		)
	`)
	if err != nil {
		panic(err)
	}

	_, err = dbConn.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS hub_user_roles (
			hub_id UUID NOT NULL,
			user_id UUID NOT NULL,
			role_id SMALLINT NOT NULL,
			PRIMARY KEY (hub_id, user_id, role_id),
			FOREIGN KEY (hub_id, user_id) REFERENCES hub_users(hub_id, user_id) ON DELETE CASCADE,
			FOREIGN KEY (hub_id, role_id) REFERENCES hub_roles(hub_id, role_id) ON DELETE CASCADE
		)
	`)
	if err != nil {
		panic(err)
	}
}

func UserSave(ctx context.Context, user *types.User) error {
	err := dbConn.QueryRow(ctx, `
        INSERT INTO users (name, pass_hash, pronouns, rgba, created_at)
        VALUES ($1, $2, $3, $4, $5)
        RETURNING id
    `, user.Name, user.PasswordHash, user.Pronouns, convertions.RgbaToUint32(user.Color), user.CreatedAt).
		Scan(&user.Id)
	return err
}

func UserDelete(ctx context.Context, id uuid.UUID) error {
	_, err := dbConn.Exec(ctx, `
		DELETE FROM users WHERE id = $1
	`, id)
	return err
}

func UserGetStandardInfoByName(ctx context.Context, user *types.User) error {
	var rgba int64
	err := dbConn.QueryRow(ctx, `
		SELECT id, name, pass_hash, COALESCE(pronouns, ''), rgba, created_at, COALESCE(avatar, ''), COALESCE(profile_bg, '') FROM users WHERE name = $1
	`, user.Name).Scan(&user.Id, &user.Name, &user.PasswordHash, &user.Pronouns, &rgba, &user.CreatedAt, &user.Avatar, &user.ProfileBg)
	if err == nil {
		user.Color = convertions.Uint32ToRgba(uint32(rgba))
	}
	return err
}

func UserGetById(ctx context.Context, user *types.User) error {
	var rgba int64
	err := dbConn.QueryRow(ctx, `
		SELECT name, pass_hash, COALESCE(pronouns, ''), rgba, created_at, COALESCE(avatar, ''), COALESCE(profile_bg, '') FROM users WHERE id = $1
	`, user.Id).Scan(&user.Name, &user.PasswordHash, &user.Pronouns, &rgba, &user.CreatedAt, &user.Avatar, &user.ProfileBg)
	if err == nil {
		user.Color = convertions.Uint32ToRgba(uint32(rgba))
	}
	return err
}

func UserUpdateProfile(ctx context.Context, user *types.User, updateList types.UserProfileUpdateList) error {
	setClauses := make([]string, 0, 3)
	args := make([]any, 0, 4)
	argIdx := 1

	if updateList.Pronouns {
		setClauses = append(setClauses, fmt.Sprintf("pronouns = $%d", argIdx))
		args = append(args, user.Pronouns)
		argIdx++
	}
	if updateList.Description {
		setClauses = append(setClauses, fmt.Sprintf("description = $%d", argIdx))
		args = append(args, user.Description)
		argIdx++
	}
	if updateList.Color {
		setClauses = append(setClauses, fmt.Sprintf("rgba = $%d", argIdx))
		args = append(args, convertions.RgbaToUint32(user.Color))
		argIdx++
	}
	if updateList.Avatar {
		setClauses = append(setClauses, fmt.Sprintf("avatar = $%d", argIdx))
		args = append(args, user.Avatar)
		argIdx++
	}
	if updateList.ProfileBg {
		setClauses = append(setClauses, fmt.Sprintf("profile_bg = $%d", argIdx))
		args = append(args, user.ProfileBg)
		argIdx++
	}

	if len(setClauses) == 0 {
		return nil
	}

	query := "UPDATE users SET " + strings.Join(setClauses, ", ") + fmt.Sprintf(" WHERE id = $%d", argIdx)
	args = append(args, user.Id)

	_, err := dbConn.Exec(ctx, query, args...)
	return err
}

func GetWholeUser(ctx context.Context, user *types.User) error {
	if err := UserGetById(ctx, user); err != nil {
		return err
	}
	if err := ConnectionsGetBelongingToUser(ctx, user); err != nil {
		return err
	}

	cache.SaveUser(user)
	return nil
}

func ConnectionSave(ctx context.Context, conn *types.Connection) error {
	return dbConn.QueryRow(ctx, `
		INSERT INTO user_connections (requestor_id, recipient_id, state, created_at) VALUES ($1, $2, $3, $4)
		RETURNING id
	`, conn.RequestorId, conn.RecipientId, conn.State, conn.CreatedAt).Scan(&conn.Id)
}

func ConnectionDelete(ctx context.Context, conn *types.Connection) error {
	_, err := dbConn.Exec(ctx, `
		DELETE FROM user_connections WHERE id = $1
	`, conn.Id)
	return err
}

func ConnectionsGetBelongingToUser(ctx context.Context, user *types.User) error {
	rows, err := dbConn.Query(ctx, `
		SELECT id, requestor_id, recipient_id, state, created_at
		FROM user_connections
		WHERE requestor_id = $1 OR recipient_id = $1
	`, user.Id)
	if err != nil {
		return err
	}
	defer rows.Close()

	if user.Connections == nil {
		user.Connections = make(map[uuid.UUID]*types.Connection)
	}

	for rows.Next() {
		conn := types.NewConnection()
		if err = rows.Scan(&conn.Id, &conn.RequestorId, &conn.RecipientId, &conn.State, &conn.CreatedAt); err != nil {
			return fmt.Errorf("scanning connection row: %w", err)
		}
		user.Connections[conn.Id] = conn
	}
	return rows.Err()
}

func ConnectionUpdateState(ctx context.Context, conn *types.Connection) error {
	_, err := dbConn.Exec(ctx, `
		UPDATE user_connections SET state = $1 WHERE id = $2
	`, conn.State, conn.Id)
	return err
}

func ConnectionMessageSave(ctx context.Context, message *types.Message) error {
	if message.Id != (uuid.UUID{}) {
		_, err := dbConn.Exec(ctx, `
			INSERT INTO direct_messages (id, sender_id, receiver_id, created_at, content, attached_file) VALUES ($1, $2, $3, $4, $5, $6)
		`, message.Id, message.Sender, message.Receiver, message.CreatedAt, message.Content, message.AttachedFile)
		return err
	}
	return dbConn.QueryRow(ctx, `
		INSERT INTO direct_messages (sender_id, receiver_id, created_at, content, attached_file) VALUES ($1, $2, $3, $4, $5)
		RETURNING id
	`, message.Sender, message.Receiver, message.CreatedAt, message.Content, message.AttachedFile).Scan(&message.Id)
}

func ConnectionGetMessagesBefore(ctx context.Context, before time.Time, connection uuid.UUID, cap uint32) ([]*types.Message, error) {
	rows, err := dbConn.Query(ctx, `
		SELECT id, sender_id, receiver_id, created_at, content, attached_file
		FROM (
			SELECT id, sender_id, receiver_id, created_at, content, attached_file
			FROM direct_messages
			WHERE receiver_id = $1
			  AND created_at < $2
			ORDER BY created_at DESC
			LIMIT $3
		) sub
		ORDER BY created_at ASC
	`, connection, before, cap)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	messages := make([]*types.Message, 0, cap)
	for rows.Next() {
		msg := &types.Message{}
		if err = rows.Scan(&msg.Id, &msg.Sender, &msg.Receiver, &msg.CreatedAt, &msg.Content, &msg.AttachedFile); err != nil {
			return nil, err
		}
		messages = append(messages, msg)
	}
	return messages, rows.Err()
}

func HubSave(ctx context.Context, hub *types.Hub) error {
	_, err := dbConn.Exec(ctx, `
		INSERT INTO hubs (id, creator_id, name, allow_user_color, rgba, created_at) VALUES ($1, $2, $3, $4, $5, $6)
	`, hub.Id, hub.Creator, hub.Name, hub.AllowUserColor, convertions.RgbaToUint32(hub.Color), hub.CreatedAt)
	return err
}

func HubDelete(ctx context.Context, hub *types.Hub) error {
	_, err := dbConn.Exec(ctx, `DELETE FROM hubs WHERE id = $1`, hub.Id)
	return err
}

func HubGet(ctx context.Context, hub *types.Hub) error {
	var rgba uint32
	err := dbConn.QueryRow(ctx, `
		SELECT creator_id, name, allow_user_color, rgba, created_at FROM hubs WHERE id = $1
	`, hub.Id).Scan(&hub.Creator, &hub.Name, &hub.AllowUserColor, &rgba, &hub.CreatedAt)
	if err == nil {
		hub.Color = convertions.Uint32ToRgba(rgba)
	}
	return err
}

func HubsGetBelongingToUser(ctx context.Context, userId uuid.UUID) ([]*types.Hub, error) {
	rows, err := dbConn.Query(ctx, `
		SELECT h.id, h.creator_id, h.name, h.allow_user_color, h.rgba, h.created_at
		FROM hubs h
		JOIN hub_users hu ON h.id = hu.hub_id
		WHERE hu.user_id = $1
	`, userId)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var hubs []*types.Hub
	for rows.Next() {
		hub := types.NewHub()
		var rgba uint32
		if err = rows.Scan(&hub.Id, &hub.Creator, &hub.Name, &hub.AllowUserColor, &rgba, &hub.CreatedAt); err != nil {
			return nil, fmt.Errorf("scanning hub row: %w", err)
		}
		hub.Color = convertions.Uint32ToRgba(rgba)
		hubs = append(hubs, hub)
	}
	return hubs, rows.Err()
}

func HubGetUsers(ctx context.Context, hub *types.Hub) error {
	rows, err := dbConn.Query(ctx, `
		SELECT hu.user_id, hu.name, hu.created_at, COALESCE(array_agg(hur.role_id) FILTER (WHERE hur.role_id IS NOT NULL), '{}')
		FROM hub_users hu
		LEFT JOIN hub_user_roles hur ON hu.hub_id = hur.hub_id AND hu.user_id = hur.user_id
		WHERE hu.hub_id = $1
		GROUP BY hu.user_id, hu.name, hu.created_at
	`, hub.Id)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		hu := &types.HubUser{}
		var roles []int16
		if err = rows.Scan(&hu.Id, &hu.Username, &hu.CreatedAt, &roles); err != nil {
			return fmt.Errorf("scanning hub user row: %w", err)
		}
		hu.Roles = make([]uint8, len(roles))
		for i, r := range roles {
			hu.Roles[i] = uint8(r)
		}
		hub.Users[hu.Id] = hu
	}
	return rows.Err()
}

func HubGetRoles(ctx context.Context, hub *types.Hub) error {
	rows, err := dbConn.Query(ctx, `
		SELECT role_id, name, permissions, rgba FROM hub_roles WHERE hub_id = $1
	`, hub.Id)
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		hr := &types.HubRole{}
		var rgba uint32
		if err = rows.Scan(&hr.Id, &hr.Name, &hr.RolePermission, &rgba); err != nil {
			return fmt.Errorf("scanning hub role row: %w", err)
		}
		hr.Color = convertions.Uint32ToRgba(rgba)
		hub.Roles[hr.Id] = hr
	}
	return rows.Err()
}

func GetWholeHub(ctx context.Context, hub *types.Hub) error {
	if err := HubGet(ctx, hub); err != nil {
		return err
	}
	if err := HubGetUsers(ctx, hub); err != nil {
		return err
	}
	return HubGetRoles(ctx, hub)
}

func HubUserAdd(ctx context.Context, hubId uuid.UUID, hu *types.HubUser) error {
	_, err := dbConn.Exec(ctx, `
		INSERT INTO hub_users (hub_id, user_id, name, created_at) VALUES ($1, $2, $3, $4)
	`, hubId, hu.Id, hu.Username, hu.CreatedAt)
	return err
}

func HubUserRemove(ctx context.Context, hubId uuid.UUID, userId uuid.UUID) error {
	_, err := dbConn.Exec(ctx, `
		DELETE FROM hub_users WHERE hub_id = $1 AND user_id = $2
	`, hubId, userId)
	return err
}

func HubUserRoleAdd(ctx context.Context, hubId uuid.UUID, userId uuid.UUID, roleId uint8) error {
	_, err := dbConn.Exec(ctx, `
		INSERT INTO hub_user_roles (hub_id, user_id, role_id) VALUES ($1, $2, $3)
	`, hubId, userId, roleId)
	return err
}

func HubUserRoleRemove(ctx context.Context, hubId uuid.UUID, userId uuid.UUID, roleId uint8) error {
	_, err := dbConn.Exec(ctx, `
		DELETE FROM hub_user_roles WHERE hub_id = $1 AND user_id = $2 AND role_id = $3
	`, hubId, userId, roleId)
	return err
}

func HubRoleSave(ctx context.Context, hubId uuid.UUID, role *types.HubRole) error {
	_, err := dbConn.Exec(ctx, `
		INSERT INTO hub_roles (hub_id, role_id, name, permissions, rgba) VALUES ($1, $2, $3, $4, $5)
	`, hubId, role.Id, role.Name, role.RolePermission, convertions.RgbaToUint32(role.Color))
	return err
}

func HubRoleDelete(ctx context.Context, hubId uuid.UUID, roleId uint8) error {
	_, err := dbConn.Exec(ctx, `
		DELETE FROM hub_roles WHERE hub_id = $1 AND role_id = $2
	`, hubId, roleId)
	return err
}

func HubRoleUpdate(ctx context.Context, hubId uuid.UUID, role *types.HubRole) error {
	_, err := dbConn.Exec(ctx, `
		UPDATE hub_roles SET name = $1, permissions = $2, rgba = $3 WHERE hub_id = $4 AND role_id = $5
	`, role.Name, role.RolePermission, convertions.RgbaToUint32(role.Color), hubId, role.Id)
	return err
}

func HubChannelGroupSave(ctx context.Context, hubId uuid.UUID, cg *types.HubChannelGroup) error {
	_, err := dbConn.Exec(ctx, `
		INSERT INTO hub_channel_groups (channel_group_id, hub_id, name, rgba, position) VALUES ($1, $2, $3, $4, $5)
	`, cg.Id, hubId, cg.Name, convertions.RgbaToUint32(cg.Color), cg.Position)
	return err
}

func HubChannelGroupDelete(ctx context.Context, cgId uuid.UUID) error {
	_, err := dbConn.Exec(ctx, `DELETE FROM hub_channel_groups WHERE channel_group_id = $1`, cgId)
	return err
}

func HubChannelSave(ctx context.Context, hubId uuid.UUID, ch *types.HubChannel) error {
	return dbConn.QueryRow(ctx, `
		INSERT INTO hub_channels (hub_id, name, parent_group_id, position) VALUES ($1, $2, $3, $4)
		RETURNING channel_id
	`, hubId, ch.Name, ch.ParentGroupId, ch.Position).Scan(&ch.Id)
}

func HubChannelDelete(ctx context.Context, chId uuid.UUID) error {
	_, err := dbConn.Exec(ctx, `DELETE FROM hub_channels WHERE channel_id = $1`, chId)
	return err
}