work on optimising vcenter queries

db/helpers.go

@@ -0,0 +1,169 @@
+package db
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"strings"
+
+	"vctp/db/queries"
+
+	"github.com/jmoiron/sqlx"
+)
+
+// SnapshotTotals summarizes counts and allocations for snapshot tables.
+type SnapshotTotals struct {
+	VmCount   int64   `db:"vm_count"`
+	VcpuTotal int64   `db:"vcpu_total"`
+	RamTotal  int64   `db:"ram_total"`
+	DiskTotal float64 `db:"disk_total"`
+}
+
+// ValidateTableName ensures table identifiers are safe for interpolation.
+func ValidateTableName(name string) error {
+	if name == "" {
+		return fmt.Errorf("table name is empty")
+	}
+	for _, r := range name {
+		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' {
+			continue
+		}
+		return fmt.Errorf("invalid table name: %s", name)
+	}
+	return nil
+}
+
+// SafeTableName returns the name if it passes validation.
+func SafeTableName(name string) (string, error) {
+	if err := ValidateTableName(name); err != nil {
+		return "", err
+	}
+	return name, nil
+}
+
+// TableHasRows returns true when a table contains at least one row.
+func TableHasRows(ctx context.Context, dbConn *sqlx.DB, table string) (bool, error) {
+	if err := ValidateTableName(table); err != nil {
+		return false, err
+	}
+	query := fmt.Sprintf(`SELECT 1 FROM %s LIMIT 1`, table)
+	var exists int
+	if err := dbConn.GetContext(ctx, &exists, query); err != nil {
+		if err == sql.ErrNoRows {
+			return false, nil
+		}
+		return false, err
+	}
+	return true, nil
+}
+
+// TableExists checks if a table exists in the current schema.
+func TableExists(ctx context.Context, dbConn *sqlx.DB, table string) bool {
+	driver := strings.ToLower(dbConn.DriverName())
+	switch driver {
+	case "sqlite":
+		q := queries.New(dbConn)
+		count, err := q.SqliteTableExists(ctx, sql.NullString{String: table, Valid: table != ""})
+		return err == nil && count > 0
+	case "pgx", "postgres":
+		var count int
+		err := dbConn.GetContext(ctx, &count, `
+SELECT COUNT(1)
+FROM pg_catalog.pg_tables
+WHERE schemaname = 'public' AND tablename = $1
+`, table)
+		return err == nil && count > 0
+	default:
+		return false
+	}
+}
+
+// ColumnExists checks if a column exists in a table.
+func ColumnExists(ctx context.Context, dbConn *sqlx.DB, tableName string, columnName string) (bool, error) {
+	driver := strings.ToLower(dbConn.DriverName())
+	switch driver {
+	case "sqlite":
+		if _, err := SafeTableName(tableName); err != nil {
+			return false, err
+		}
+		query := fmt.Sprintf(`PRAGMA table_info("%s")`, tableName)
+		rows, err := dbConn.QueryxContext(ctx, query)
+		if err != nil {
+			return false, err
+		}
+		defer rows.Close()
+		for rows.Next() {
+			var (
+				cid        int
+				name       string
+				colType    string
+				notNull    int
+				defaultVal sql.NullString
+				pk         int
+			)
+			if err := rows.Scan(&cid, &name, &colType, &notNull, &defaultVal, &pk); err != nil {
+				return false, err
+			}
+			if strings.EqualFold(name, columnName) {
+				return true, nil
+			}
+		}
+		return false, rows.Err()
+	case "pgx", "postgres":
+		var count int
+		err := dbConn.GetContext(ctx, &count, `
+SELECT COUNT(1)
+FROM information_schema.columns
+WHERE table_name = $1 AND column_name = $2
+`, tableName, strings.ToLower(columnName))
+		if err != nil {
+			return false, err
+		}
+		return count > 0, nil
+	default:
+		return false, fmt.Errorf("unsupported driver for column lookup: %s", driver)
+	}
+}
+
+// SnapshotTotalsForTable returns totals for a snapshot table.
+func SnapshotTotalsForTable(ctx context.Context, dbConn *sqlx.DB, table string) (SnapshotTotals, error) {
+	if _, err := SafeTableName(table); err != nil {
+		return SnapshotTotals{}, err
+	}
+	query := fmt.Sprintf(`
+SELECT
+  COUNT(DISTINCT "VmId") AS vm_count,
+  COALESCE(SUM(CASE WHEN "VcpuCount" IS NOT NULL THEN "VcpuCount" ELSE 0 END), 0) AS vcpu_total,
+  COALESCE(SUM(CASE WHEN "RamGB" IS NOT NULL THEN "RamGB" ELSE 0 END), 0) AS ram_total,
+  COALESCE(SUM(CASE WHEN "ProvisionedDisk" IS NOT NULL THEN "ProvisionedDisk" ELSE 0 END), 0) AS disk_total
+FROM %s
+WHERE "IsPresent" = 'TRUE'
+`, table)
+
+	var totals SnapshotTotals
+	if err := dbConn.GetContext(ctx, &totals, query); err != nil {
+		return SnapshotTotals{}, err
+	}
+	return totals, nil
+}
+
+// SnapshotTotalsForUnion returns totals for a union query of snapshots.
+func SnapshotTotalsForUnion(ctx context.Context, dbConn *sqlx.DB, unionQuery string) (SnapshotTotals, error) {
+	query := fmt.Sprintf(`
+SELECT
+  COUNT(DISTINCT "VmId") AS vm_count,
+  COALESCE(SUM(CASE WHEN "VcpuCount" IS NOT NULL THEN "VcpuCount" ELSE 0 END), 0) AS vcpu_total,
+  COALESCE(SUM(CASE WHEN "RamGB" IS NOT NULL THEN "RamGB" ELSE 0 END), 0) AS ram_total,
+  COALESCE(SUM(CASE WHEN "ProvisionedDisk" IS NOT NULL THEN "ProvisionedDisk" ELSE 0 END), 0) AS disk_total
+FROM (
+%s
+) snapshots
+WHERE "IsPresent" = 'TRUE'
+`, unionQuery)
+
+	var totals SnapshotTotals
+	if err := dbConn.GetContext(ctx, &totals, query); err != nil {
+		return SnapshotTotals{}, err
+	}
+	return totals, nil
+}

db/queries/models.go

@@ -59,6 +59,15 @@ type InventoryHistory struct {
 	PreviousProvisionedDisk sql.NullFloat64 `db:"PreviousProvisionedDisk" json:"PreviousProvisionedDisk"`
 }
 
+type PragmaTableInfo struct {
+	Cid       sql.NullInt64  `db:"cid" json:"cid"`
+	Name      sql.NullString `db:"name" json:"name"`
+	Type      sql.NullString `db:"type" json:"type"`
+	Notnull   sql.NullInt64  `db:"notnull" json:"notnull"`
+	DfltValue sql.NullString `db:"dflt_value" json:"dflt_value"`
+	Pk        sql.NullInt64  `db:"pk" json:"pk"`
+}
+
 type SnapshotRegistry struct {
 	ID           int64  `db:"id" json:"id"`
 	SnapshotType string `db:"snapshot_type" json:"snapshot_type"`
@@ -66,6 +75,14 @@ type SnapshotRegistry struct {
 	SnapshotTime int64  `db:"snapshot_time" json:"snapshot_time"`
 }
 
+type SqliteMaster struct {
+	Type     sql.NullString `db:"type" json:"type"`
+	Name     sql.NullString `db:"name" json:"name"`
+	TblName  sql.NullString `db:"tbl_name" json:"tbl_name"`
+	Rootpage sql.NullInt64  `db:"rootpage" json:"rootpage"`
+	Sql      sql.NullString `db:"sql" json:"sql"`
+}
+
 type Update struct {
 	Uid                int64           `db:"Uid" json:"Uid"`
 	InventoryId        sql.NullInt64   `db:"InventoryId" json:"InventoryId"`

db/queries/query.sql

@@ -119,3 +119,13 @@ INSERT INTO inventory_history (
     ?, ?, ?, ?, ?, ?, ?
 )
 RETURNING *;
+
+-- name: SqliteTableExists :one
+SELECT COUNT(1) AS count
+FROM sqlite_master
+WHERE type = 'table' AND name = sqlc.arg('table_name');
+
+-- name: SqliteColumnExists :one
+SELECT COUNT(1) AS count
+FROM pragma_table_info
+WHERE name = sqlc.arg('column_name');

db/queries/query.sql.go

@@ -876,6 +876,32 @@ func (q *Queries) ListUnprocessedEvents(ctx context.Context, eventtime sql.NullI
 	return items, nil
 }
 
+const sqliteColumnExists = `-- name: SqliteColumnExists :one
+SELECT COUNT(1) AS count
+FROM pragma_table_info
+WHERE name = ?1
+`
+
+func (q *Queries) SqliteColumnExists(ctx context.Context, columnName sql.NullString) (int64, error) {
+	row := q.db.QueryRowContext(ctx, sqliteColumnExists, columnName)
+	var count int64
+	err := row.Scan(&count)
+	return count, err
+}
+
+const sqliteTableExists = `-- name: SqliteTableExists :one
+SELECT COUNT(1) AS count
+FROM sqlite_master
+WHERE type = 'table' AND name = ?1
+`
+
+func (q *Queries) SqliteTableExists(ctx context.Context, tableName sql.NullString) (int64, error) {
+	row := q.db.QueryRowContext(ctx, sqliteTableExists, tableName)
+	var count int64
+	err := row.Scan(&count)
+	return count, err
+}
+
 const updateEventsProcessed = `-- name: UpdateEventsProcessed :exec
 UPDATE events
 SET "Processed" = 1

db/schema.sql

@@ -72,3 +72,22 @@ CREATE TABLE IF NOT EXISTS snapshot_registry (
 	"table_name" TEXT NOT NULL UNIQUE,
 	"snapshot_time" INTEGER NOT NULL
 );
+
+-- The following tables are declared for sqlc type-checking only.
+-- Do not apply this file as a migration.
+CREATE TABLE sqlite_master (
+	"type" TEXT,
+	"name" TEXT,
+	"tbl_name" TEXT,
+	"rootpage" INTEGER,
+	"sql" TEXT
+);
+
+CREATE TABLE pragma_table_info (
+	"cid" INTEGER,
+	"name" TEXT,
+	"type" TEXT,
+	"notnull" INTEGER,
+	"dflt_value" TEXT,
+	"pk" INTEGER
+);