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) } } // DbSaveClientWithoutGroups saves client in db without groups and sets its id // return: error if not successful 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 } // DbSetClientByName sets all fields of given struct with database's data using name // return: error if not successful 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) } // DbSetClientByIdWithoutGroups sets all fields of given struct with database's data using id, excluding groups // return: error if not successful 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 } // DbSetClientById sets all fields of given struct with database's data using id, including groups // return: error if not successful func DbSetClientById(ctx context.Context, client *Client) error { err := DbSetClientByIdWithoutGroups(ctx, client) if err != nil { return err } return DbSetClientGroups(ctx, client) } // DbSaveGroupWithoutClients saves group in db and sets its id, also adds the owner as first member // return: error if not successful 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 } // DbDeleteGroup deletes given group by id // return: error if not successful func DbDeleteGroup(ctx context.Context, group *Group) error { _, err := dbConn.Exec(ctx, ` DELETE FROM chat_groups WHERE id = $1 `, group.Id) return err } // DbSetGroupByIdWithoutClients sets all fields of given struct with database's data using id, populates Clients map with member ids but not their data // return: error if not successful 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() } // DbSetGroupById sets all fields of given struct with database's data using id, including full member client data // return: error if not successful func DbSetGroupById(ctx context.Context, group *Group) error { err := DbSetGroupByIdWithoutClients(ctx, group) if err != nil { return err } return DbSetGroupMemberClients(ctx, group) } // DbSetGroupMemberClients populates group's Clients map with ids of all members from database // return: error if not successful 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() } // DbAddClientsToGroup adds given clients to group in db, silently ignores already existing members // return: error if not successful func DbAddClientsToGroup(ctx context.Context, groupId uint32, clientIds *[MaxClientsInGroup]uint32) error { batch := &pgx.Batch{} now := time.Now() var count int for _, cid := range clientIds { if cid == 0 { continue } batch.Queue(` INSERT INTO chat_group_members (group_id, user_id, joined_at) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING `, groupId, cid, now) count++ } br := dbConn.SendBatch(ctx, batch) defer br.Close() for range count { if _, err := br.Exec(); err != nil { return err } } return nil } // DbRemoveClientsFromGroup removes given clients from group in db, silently ignores not existing members // return: deleted clients count, error if not successful func DbRemoveClientsFromGroup(ctx context.Context, groupId uint32, clientIds *[MaxClientsInGroup]uint32) (int, error) { batch := &pgx.Batch{} var count int for _, cid := range clientIds { if cid == 0 { continue } batch.Queue(` DELETE FROM chat_group_members WHERE group_id = $1 AND user_id = $2 `, groupId, cid) 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 } // DbSetClientGroups populates client's Groups map with ids of all groups the client belongs to from database // return: error if not successful 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() } // DbSetGroupColor set group's color based on id // return: error if not successful func DbSetGroupColor(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 } // DbSetGroupOwnerId set group's owner based on id // return: error if not successful func DbSetGroupOwnerId(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 }