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"
)

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 {
	db              *sql.DB
	queryCuboidStmt *sql.Stmt
}

func NewPGClient(connS string) (*PGClient, error) {
	db, err := sql.Open("pgx", connS)
	if err != nil {
		return nil, err
	}
	stmt, err := db.Prepare(queryCuboidSQL)
	if err != nil {
		return nil, err
	}
	client := PGClient{
		db:              db,
		queryCuboidStmt: stmt,
	}
	return &client, 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.db.Close()
}