mtsatellite/common/pgclient.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()
}