36 lines
776 B
Go
36 lines
776 B
Go
|
package postgres
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"log/slog"
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"git.optclblast.xyz/draincloud/draincloud-core/internal/logger"
|
||
|
|
"github.com/jackc/pgx/v5"
|
||
|
|
)
|
||
|
|
|
||
|
|
type ShardMap = map[uint16]*pgx.ConnConfig
|
||
|
|
|
||
|
|
type ShardCluster struct {
|
||
|
|
m sync.Mutex
|
||
|
|
shards map[uint16]*pgx.Conn
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewShardCluster(ctx context.Context, shardMap ShardMap) *ShardCluster {
|
||
|
|
shards := make(map[uint16]*pgx.Conn, len(shardMap))
|
||
|
|
for n, cfg := range shardMap {
|
||
|
|
conn, err := pgx.ConnectConfig(ctx, cfg)
|
||
|
|
if err != nil {
|
||
|
|
logger.Fatal(ctx, "failed to connect to shard", slog.Uint64("num", uint64(n)), logger.Err(err))
|
||
|
|
}
|
||
|
|
shards[n] = conn
|
||
|
|
}
|
||
|
|
return &ShardCluster{shards: shards}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *ShardCluster) SelectShard(n uint16) *pgx.Conn {
|
||
|
|
c.m.Lock()
|
||
|
|
defer c.m.Unlock()
|
||
|
|
return c.shards[n]
|
||
|
|
}
|