mtsatellite/common/pgclient.go
2024-01-06 16:12:19 +01:00

118 lines
2.1 KiB
Go

// Copyright 2022 by Sascha L. Teichmann
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.
package common
import (
"context"
"database/sql"
_ "github.com/jackc/pgx/v4/stdlib" // link PostgreSQL driver
)
const queryCuboidSQL = `
SELECT posx, posy, posz, data FROM blocks
WHERE
posx BETWEEN $1 AND $2 AND
posy BETWEEN $3 AND $4 AND
posz BETWEEN $5 AND $6`
type PGClient struct {
conn *sql.Conn
queryCuboidStmt *sql.Stmt
}
type PGClientFactory struct {
db *sql.DB
}
func NewPGClientFactory(connS string) (*PGClientFactory, error) {
db, err := sql.Open("pgx", connS)
if err != nil {
return nil, err
}
return &PGClientFactory{db: db}, nil
}
func (pgcf *PGClientFactory) Close() error {
return pgcf.db.Close()
}
func (pgcf *PGClientFactory) Create() (DBClient, error) {
ctx := context.Background()
conn, err := pgcf.db.Conn(ctx)
if err != nil {
return nil, err
}
stmt, err := conn.PrepareContext(ctx, queryCuboidSQL)
if err != nil {
conn.Close()
return nil, err
}
return &PGClient{
conn: conn,
queryCuboidStmt: stmt,
}, nil
}
func (pgc *PGClient) QueryCuboid(
cuboid Cuboid,
fn func(*Block) *Block,
) (int, error) {
rows, err := pgc.queryCuboidStmt.QueryContext(
context.Background(),
cuboid.P1.X, cuboid.P2.X,
cuboid.P1.Y, cuboid.P2.Y,
cuboid.P1.Z, cuboid.P2.Z)
if err != nil {
return 0, err
}
defer rows.Close()
var (
posX, posY, posZ int
data []byte
count int
block *Block
)
for ; rows.Next(); count++ {
if err := rows.Scan(
&posX, &posY, &posZ,
&data,
); err != nil {
return count, err
}
c := Coord{
X: int16(posX),
Y: int16(posY),
Z: int16(posZ),
}
if block == nil {
block = &Block{Coord: c, Data: data}
} else {
*block = Block{Coord: c, Data: data}
}
if block = fn(block); block != nil {
data = block.Data[:0]
} else {
data = nil
}
}
return count, rows.Err()
}
func (pgc *PGClient) Close() error {
if pgc.queryCuboidStmt != nil {
pgc.queryCuboidStmt.Close()
}
return pgc.conn.Close()
}