208 lines
6.6 KiB
Go
208 lines
6.6 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"database/sql"
|
|
"embed"
|
|
"errors"
|
|
"fmt"
|
|
"io/fs"
|
|
"sync"
|
|
|
|
//"github.com/georgysavva/scany/v2/pgxscan"
|
|
//"github.com/jackc/pgx/v5"
|
|
"github.com/Gleipnir-Technology/nidus-sync/enums"
|
|
"github.com/Gleipnir-Technology/nidus-sync/models"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
"github.com/jackc/pgx/v5/stdlib"
|
|
_ "github.com/jackc/pgx/v5/stdlib"
|
|
"github.com/pressly/goose/v3"
|
|
"github.com/rs/zerolog/log"
|
|
"github.com/stephenafamo/bob"
|
|
"github.com/stephenafamo/bob/dialect/psql"
|
|
"github.com/stephenafamo/bob/dialect/psql/dialect"
|
|
"github.com/stephenafamo/bob/dialect/psql/im"
|
|
"github.com/uber/h3-go/v4"
|
|
)
|
|
|
|
//go:embed migrations/*.sql
|
|
var embedMigrations embed.FS
|
|
|
|
type postgres struct {
|
|
BobDB bob.DB
|
|
PGXPool *pgxpool.Pool
|
|
}
|
|
|
|
var (
|
|
PGInstance *postgres
|
|
pgOnce sync.Once
|
|
)
|
|
|
|
func doMigrations(connection_string string) error {
|
|
log.Info().Str("dsn", connection_string).Msg("Connecting to database")
|
|
db, err := sql.Open("pgx", connection_string)
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to open database connection: %w", err)
|
|
}
|
|
defer db.Close()
|
|
row := db.QueryRowContext(context.Background(), "SELECT version()")
|
|
var val string
|
|
if err := row.Scan(&val); err != nil {
|
|
return fmt.Errorf("Failed to get database version query result: %w", err)
|
|
}
|
|
log.Info().Str("version", val).Msg("Connected to database")
|
|
|
|
fsys, err := fs.Sub(embedMigrations, "migrations")
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to get migrations embedded directory: %w", err)
|
|
}
|
|
provider, err := goose.NewProvider(goose.DialectPostgres, db, fsys)
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to create goose provider: %w", err)
|
|
}
|
|
//goose.SetBaseFS(embedMigrations)
|
|
|
|
current, target, err := provider.GetVersions(context.Background())
|
|
if err != nil {
|
|
return fmt.Errorf("Faield to get goose versions: %w", err)
|
|
}
|
|
log.Info().Int("current", int(current)).Int("target", int(target)).Msg("Migration status")
|
|
results, err := provider.Up(context.Background())
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to run migrations: %w", err)
|
|
}
|
|
if len(results) > 0 {
|
|
for _, r := range results {
|
|
log.Info().Int("version", int(r.Source.Version)).Str("direction", r.Direction).Msg("Migration done")
|
|
}
|
|
} else {
|
|
log.Info().Msg("No migrations necessary.")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func initializeDatabase(ctx context.Context, uri string) error {
|
|
needs, err := needsMigrations(uri)
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to determine if migrations are needed: %w", err)
|
|
}
|
|
if needs == nil {
|
|
return errors.New("Can't read variable 'needs' - it's nil")
|
|
}
|
|
if *needs {
|
|
//return errors.New(fmt.Sprintf("Must migrate database before connecting: %t", *needs))
|
|
log.Info().Msg("Handling database migrations")
|
|
err = doMigrations(uri)
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to handle migrations: %w", err)
|
|
}
|
|
} else {
|
|
log.Info().Msg("No database migrations necessary")
|
|
}
|
|
|
|
pgOnce.Do(func() {
|
|
db, e := pgxpool.New(ctx, uri)
|
|
bobDB := bob.NewDB(stdlib.OpenDBFromPool(db))
|
|
PGInstance = &postgres{bobDB, db}
|
|
err = e
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("unable to create connection pool: %w", err)
|
|
}
|
|
|
|
var current string
|
|
query := `SELECT current_database()`
|
|
err = PGInstance.BobDB.QueryRow(query).Scan(¤t)
|
|
if err != nil {
|
|
return fmt.Errorf("Failed to get database current: %w", err)
|
|
}
|
|
log.Info().Str("database", current).Msg("Connected to database")
|
|
return nil
|
|
}
|
|
|
|
func needsMigrations(connection_string string) (*bool, error) {
|
|
log.Info().Str("dsn", connection_string).Msg("Connecting to database")
|
|
db, err := sql.Open("pgx", connection_string)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Failed to open database connection: %w", err)
|
|
}
|
|
defer db.Close()
|
|
row := db.QueryRowContext(context.Background(), "SELECT version()")
|
|
var val string
|
|
if err := row.Scan(&val); err != nil {
|
|
return nil, fmt.Errorf("Failed to get database version query result: %w", err)
|
|
}
|
|
log.Info().Str("dsn", val).Msg("Connected to database")
|
|
|
|
fsys, err := fs.Sub(embedMigrations, "migrations")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Failed to get migrations embedded directory: %w", err)
|
|
}
|
|
provider, err := goose.NewProvider(goose.DialectPostgres, db, fsys)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Failed to create goose provider: %w", err)
|
|
}
|
|
|
|
hasPending, err := provider.HasPending(context.Background())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &hasPending, nil
|
|
}
|
|
|
|
func updateSummaryTables(ctx context.Context, org *models.Organization) {
|
|
/*org, err := models.FindOrganization(ctx, PGInstance.BobDB, org_id)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("Failed to get organization")
|
|
}*/
|
|
log.Info().Int("org_id", int(org.ID)).Msg("Getting point locations")
|
|
point_locations, err := org.FSPointlocations().All(ctx, PGInstance.BobDB)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("Failed to get organization")
|
|
return
|
|
}
|
|
log.Info().Int("count", len(point_locations)).Msg("Summarizing point locations")
|
|
|
|
for i := range 16 {
|
|
log.Info().Int("resolution", i).Msg("Working summary layer")
|
|
cellToCount := make(map[h3.Cell]int, 0)
|
|
for _, p := range point_locations {
|
|
cell, err := getCell(p.GeometryX, p.GeometryY, i)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("Failed to get cell")
|
|
continue
|
|
}
|
|
//log.Info().Float64("X", p.GeometryX).Float64("Y", p.GeometryY).Str("cell", cell.String()).Msg("Converted lat/lng")
|
|
cellToCount[cell] = cellToCount[cell] + 1
|
|
}
|
|
var to_insert []bob.Mod[*dialect.InsertQuery] = make([]bob.Mod[*dialect.InsertQuery], 0)
|
|
to_insert = append(to_insert, im.Into("h3_aggregation", "cell", "resolution", "count_", "type_", "organization_id", "geometry"))
|
|
for cell, count := range cellToCount {
|
|
polygon, err := cellToPostgisGeometry(cell)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("Failed to get PostGIS geometry")
|
|
continue
|
|
}
|
|
// log.Info().Str("polygon", polygon).Msg("Going to insert")
|
|
to_insert = append(to_insert, im.Values(psql.Arg(cell.String(), i, count, enums.H3aggregationtypeServicerequest, org.ID), psql.F("st_geomfromtext", psql.S(polygon), 4326)))
|
|
}
|
|
to_insert = append(to_insert, im.OnConflict("cell, organization_id, type_").DoUpdate(
|
|
im.SetCol("count_").To(psql.Raw("EXCLUDED.count_")),
|
|
))
|
|
//log.Info().Str("sql", insertQueryToString(psql.Insert(to_insert...))).Msg("Updating...")
|
|
_, err := psql.Insert(to_insert...).Exec(ctx, PGInstance.BobDB)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("Faild to add h3 aggregation")
|
|
}
|
|
}
|
|
}
|
|
|
|
func insertQueryToString(query bob.BaseQuery[*dialect.InsertQuery]) string {
|
|
buf := new(bytes.Buffer)
|
|
_, err := query.WriteQuery(context.TODO(), buf, 0)
|
|
if err != nil {
|
|
return fmt.Sprintf("Failed to write query: %v", err)
|
|
}
|
|
return buf.String()
|
|
}
|