Add PostgreSQL, user/squad management, remove private domains from docs

2026-05-21 01:13:23 +03:00
parent d0dc8d5822
commit 5e3229e998
17 changed files with 1171 additions and 58 deletions
@@ -0,0 +1,65 @@
+package db
+
+import (
+	"context"
+	"embed"
+	"fmt"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+//go:embed migrations/*.sql
+var migrationsFS embed.FS
+
+type DB struct {
+	pool *pgxpool.Pool
+}
+
+func Connect(ctx context.Context, databaseURL string) (*DB, error) {
+	cfg, err := pgxpool.ParseConfig(databaseURL)
+	if err != nil {
+		return nil, fmt.Errorf("parse database url: %w", err)
+	}
+	cfg.MaxConns = 10
+	cfg.MinConns = 1
+
+	pool, err := pgxpool.NewWithConfig(ctx, cfg)
+	if err != nil {
+		return nil, fmt.Errorf("connect postgres: %w", err)
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	if err := pool.Ping(ctx); err != nil {
+		pool.Close()
+		return nil, fmt.Errorf("ping postgres: %w", err)
+	}
+
+	d := &DB{pool: pool}
+	if err := d.migrate(ctx); err != nil {
+		pool.Close()
+		return nil, err
+	}
+	return d, nil
+}
+
+func (d *DB) Close() {
+	d.pool.Close()
+}
+
+func (d *DB) migrate(ctx context.Context) error {
+	data, err := migrationsFS.ReadFile("migrations/001_init.sql")
+	if err != nil {
+		return fmt.Errorf("read migration: %w", err)
+	}
+	_, err = d.pool.Exec(ctx, string(data))
+	if err != nil {
+		return fmt.Errorf("apply migration: %w", err)
+	}
+	return nil
+}
+
+func (d *DB) Pool() *pgxpool.Pool {
+	return d.pool
+}
@@ -0,0 +1,29 @@
+CREATE TABLE IF NOT EXISTS telegram_users (
+    id BIGSERIAL PRIMARY KEY,
+    telegram_id BIGINT NOT NULL UNIQUE,
+    username TEXT,
+    first_name TEXT,
+    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+);
+
+CREATE TABLE IF NOT EXISTS vpn_users (
+    id BIGSERIAL PRIMARY KEY,
+    telegram_id BIGINT,
+    remnawave_uuid UUID NOT NULL UNIQUE,
+    remnawave_username VARCHAR(36) NOT NULL,
+    external_squad_uuid UUID,
+    internal_squad_uuids UUID[] NOT NULL DEFAULT '{}',
+    expire_at TIMESTAMPTZ,
+    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+);
+
+CREATE INDEX IF NOT EXISTS idx_vpn_users_telegram ON vpn_users(telegram_id);
+CREATE INDEX IF NOT EXISTS idx_vpn_users_username ON vpn_users(remnawave_username);
+
+CREATE TABLE IF NOT EXISTS admin_wizard (
+    admin_telegram_id BIGINT PRIMARY KEY,
+    step TEXT NOT NULL,
+    data JSONB NOT NULL DEFAULT '{}',
+    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+);
@@ -0,0 +1,75 @@
+package db
+
+import (
+	"context"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+)
+
+type VPNUser struct {
+	ID                  int64
+	TelegramID          *int64
+	RemnawaveUUID       string
+	RemnawaveUsername   string
+	ExternalSquadUUID   *string
+	InternalSquadUUIDs  []string
+	ExpireAt            *time.Time
+}
+
+func (d *DB) UpsertTelegramUser(ctx context.Context, telegramID int64, username, firstName string) error {
+	_, err := d.pool.Exec(ctx, `
+		INSERT INTO telegram_users (telegram_id, username, first_name)
+		VALUES ($1, $2, $3)
+		ON CONFLICT (telegram_id) DO UPDATE SET
+			username = EXCLUDED.username,
+			first_name = EXCLUDED.first_name`,
+		telegramID, nullStr(username), nullStr(firstName))
+	return err
+}
+
+func (d *DB) SaveVPNUser(ctx context.Context, u VPNUser) error {
+	_, err := d.pool.Exec(ctx, `
+		INSERT INTO vpn_users (
+			telegram_id, remnawave_uuid, remnawave_username,
+			external_squad_uuid, internal_squad_uuids, expire_at
+		) VALUES ($1, $2, $3, $4, $5, $6)
+		ON CONFLICT (remnawave_uuid) DO UPDATE SET
+			telegram_id = EXCLUDED.telegram_id,
+			remnawave_username = EXCLUDED.remnawave_username,
+			external_squad_uuid = EXCLUDED.external_squad_uuid,
+			internal_squad_uuids = EXCLUDED.internal_squad_uuids,
+			expire_at = EXCLUDED.expire_at,
+			updated_at = NOW()`,
+		u.TelegramID, u.RemnawaveUUID, u.RemnawaveUsername,
+		u.ExternalSquadUUID, u.InternalSquadUUIDs, u.ExpireAt)
+	return err
+}
+
+func (d *DB) GetVPNByUsername(ctx context.Context, username string) (*VPNUser, error) {
+	row := d.pool.QueryRow(ctx, `
+		SELECT id, telegram_id, remnawave_uuid::text, remnawave_username,
+			external_squad_uuid::text, internal_squad_uuids::text[], expire_at
+		FROM vpn_users WHERE remnawave_username = $1`, username)
+
+	var u VPNUser
+	var ext *string
+	var internal []string
+	err := row.Scan(&u.ID, &u.TelegramID, &u.RemnawaveUUID, &u.RemnawaveUsername, &ext, &internal, &u.ExpireAt)
+	if err == pgx.ErrNoRows {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, err
+	}
+	u.ExternalSquadUUID = ext
+	u.InternalSquadUUIDs = internal
+	return &u, nil
+}
+
+func nullStr(s string) *string {
+	if s == "" {
+		return nil
+	}
+	return &s
+}
@@ -0,0 +1,123 @@
+package db
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+)
+
+type WizardData map[string]any
+
+type AdminWizard struct {
+	AdminID int64
+	Step    string
+	Data    WizardData
+}
+
+func (d *DB) GetWizard(ctx context.Context, adminID int64) (*AdminWizard, error) {
+	row := d.pool.QueryRow(ctx, `
+		SELECT admin_telegram_id, step, data
+		FROM admin_wizard WHERE admin_telegram_id = $1`, adminID)
+
+	var w AdminWizard
+	var raw []byte
+	if err := row.Scan(&w.AdminID, &w.Step, &raw); err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil, nil
+		}
+		return nil, err
+	}
+	if len(raw) > 0 {
+		_ = json.Unmarshal(raw, &w.Data)
+	}
+	if w.Data == nil {
+		w.Data = WizardData{}
+	}
+	return &w, nil
+}
+
+func (d *DB) SetWizard(ctx context.Context, adminID int64, step string, data WizardData) error {
+	raw, _ := json.Marshal(data)
+	_, err := d.pool.Exec(ctx, `
+		INSERT INTO admin_wizard (admin_telegram_id, step, data, updated_at)
+		VALUES ($1, $2, $3, NOW())
+		ON CONFLICT (admin_telegram_id) DO UPDATE SET
+			step = EXCLUDED.step,
+			data = EXCLUDED.data,
+			updated_at = NOW()`,
+		adminID, step, raw)
+	return err
+}
+
+func (d *DB) ClearWizard(ctx context.Context, adminID int64) error {
+	_, err := d.pool.Exec(ctx, `DELETE FROM admin_wizard WHERE admin_telegram_id = $1`, adminID)
+	return err
+}
+
+func (w WizardData) String(key string) string {
+	v, _ := w[key].(string)
+	return v
+}
+
+func (w WizardData) Int(key string) int {
+	switch v := w[key].(type) {
+	case float64:
+		return int(v)
+	case int:
+		return v
+	default:
+		return 0
+	}
+}
+
+func (w WizardData) StringSlice(key string) []string {
+	raw, ok := w[key].([]any)
+	if !ok {
+		if ss, ok := w[key].([]string); ok {
+			return ss
+		}
+		return nil
+	}
+	out := make([]string, 0, len(raw))
+	for _, x := range raw {
+		if s, ok := x.(string); ok {
+			out = append(out, s)
+		}
+	}
+	return out
+}
+
+func (w WizardData) Set(key string, val any) {
+	w[key] = val
+}
+
+func (w WizardData) ToggleUUID(key, uuid string) {
+	cur := w.StringSlice(key)
+	for i, id := range cur {
+		if id == uuid {
+			cur = append(cur[:i], cur[i+1:]...)
+			w[key] = cur
+			return
+		}
+	}
+	w[key] = append(cur, uuid)
+}
+
+const (
+	StepIdle              = ""
+	StepAwaitUsername     = "await_username"
+	StepAwaitDays         = "await_days"
+	StepPickExternalSquad = "pick_external"
+	StepPickInternalSquads = "pick_internal"
+	StepConfirm           = "confirm"
+)
+
+func DefaultExpireAt(days int) time.Time {
+	if days <= 0 {
+		days = 30
+	}
+	return time.Now().UTC().AddDate(0, 0, days)
+}