data/postgres.go - master - Magister/Venera_20

data/postgres.go
275 строк8 KB
Magister
02 июл 2026, 23:59
02 июл 2026, 23:59182c551
package data

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	// PostgreSQL driver
	_ "github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// PostgreSQL - интерфейс для работы с PostgreSQL
type PostgreSQL struct {
	pool *pgxpool.Pool
}

// PostgreSQLConfig - конфигурация подключения к PostgreSQL
type PostgreSQLConfig struct {
	Host       string
	Port       int
	Database   string
	Username   string
	Password   string
}

// NewPostgreSQL - создает новый клиент PostgreSQL
func NewPostgreSQL(config PostgreSQLConfig) (*PostgreSQL, error) {
	connStr := fmt.Sprintf("postgresql://%s:%s@%s:%d/%s",
		config.Username, config.Password, config.Host, config.Port, config.Database)

	pool, err := pgxpool.New(context.Background(), connStr)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to PostgreSQL: %w", err)
	}

	// Проверка подключения
	err = pool.Ping(context.Background())
	if err != nil {
		pool.Close()
		return nil, fmt.Errorf("failed to ping PostgreSQL: %w", err)
	}

	return &PostgreSQL{
		pool: pool,
	}, nil
}

// Close - закрывает соединение с PostgreSQL
func (db *PostgreSQL) Close() {
	db.pool.Close()
}

// Ping - проверка подключения к PostgreSQL
func (db *PostgreSQL) Ping() error {
	return db.pool.Ping(context.Background())
}

// InsertData - вставляет данные в таблицу
func (db *PostgreSQL) InsertData(sourceName, key, value string) (int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var id int64
	err := db.pool.QueryRow(ctx,
		`INSERT INTO data_records (source_name, key, value, date_first, date_last)
		 VALUES ($1, $2, $3, NOW(), NOW())
		 ON CONFLICT (source_name, key, value)
		 DO UPDATE SET date_last = NOW()
		 RETURNING id`,
		sourceName, key, value).Scan(&id)

	if err != nil {
		return 0, fmt.Errorf("failed to insert data: %w", err)
	}

	return id, nil
}

// InsertBatchData - вставляет пакет данных
func (db *PostgreSQL) InsertBatchData(records []DataRecord) error {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	// Используем COPY для быстрой вставки
	_, err := db.pool.CopyFrom(ctx,
		pgx.Identifier{"data_records"},
		[]string{"source_name", "key", "value", "date_first", "date_last"},
		pgx.CopyFromRows(records))

	if err != nil {
		return fmt.Errorf("failed to insert batch data: %w", err)
	}

	return nil
}

// GetData - получает данные по фильтру
func (db *PostgreSQL) GetData(sourceName, key, value string, dateFrom, dateTo int64, page, pageSize int) ([]DataRecord, int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Формирование запроса
	query := `SELECT id, source_name, key, value, date_first, date_last FROM data_records WHERE 1=1`
	countQuery := `SELECT COUNT(*) FROM data_records WHERE 1=1`

	args := []interface{}{}
	argIndex := 1

	if sourceName != "" {
		query += fmt.Sprintf(" AND source_name = $%d", argIndex)
		countQuery += fmt.Sprintf(" AND source_name = $%d", argIndex)
		args = append(args, sourceName)
		argIndex++
	}

	if key != "" {
		query += fmt.Sprintf(" AND key = $%d", argIndex)
		countQuery += fmt.Sprintf(" AND key = $%d", argIndex)
		args = append(args, key)
		argIndex++
	}

	if value != "" {
		query += fmt.Sprintf(" AND value = $%d", argIndex)
		countQuery += fmt.Sprintf(" AND value = $%d", argIndex)
		args = append(args, value)
		argIndex++
	}

	if dateFrom > 0 {
		query += fmt.Sprintf(" AND date_first >= to_timestamp($%d)", argIndex)
		countQuery += fmt.Sprintf(" AND date_first >= to_timestamp($%d)", argIndex)
		args = append(args, time.Unix(dateFrom, 0))
		argIndex++
	}

	if dateTo > 0 {
		query += fmt.Sprintf(" AND date_first <= to_timestamp($%d)", argIndex)
		countQuery += fmt.Sprintf(" AND date_first <= to_timestamp($%d)", argIndex)
		args = append(args, time.Unix(dateTo, 0))
		argIndex++
	}

	// Получение количества записей
	var total int64
	err := db.pool.QueryRow(ctx, countQuery, args...).Scan(&total)
	if err != nil {
		return nil, 0, fmt.Errorf("failed to get count: %w", err)
	}

	// Добавление пагинации
	query += " ORDER BY date_first DESC LIMIT $1 OFFSET $2"
	args = append(args, pageSize, (page-1)*pageSize)

	// Выполнение запроса
	rows, err := db.pool.Query(ctx, query, args...)
	if err != nil {
		return nil, 0, fmt.Errorf("failed to query data: %w", err)
	}
	defer rows.Close()

	records := make([]DataRecord, 0)
	for rows.Next() {
		var record DataRecord
		err := rows.Scan(&record.ID, &record.SourceName, &record.Key, &record.Value, &record.FirstAppearance, &record.LastAppearance)
		if err != nil {
			return nil, 0, fmt.Errorf("failed to scan record: %w", err)
		}
		records = append(records, record)
	}

	return records, total, nil
}

// GetStatistics - получает статистику
func (db *PostgreSQL) GetStatistics() (map[string]interface{}, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	stats := make(map[string]interface{})

	// Общее количество записей
	var totalRecords int64
	err := db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM data_records`).Scan(&totalRecords)
	if err != nil {
		return nil, fmt.Errorf("failed to get total records: %w", err)
	}
	stats["total_records"] = totalRecords

	// Количество уникальных источников
	var uniqueSources int64
	err = db.pool.QueryRow(ctx, `SELECT COUNT(DISTINCT source_name) FROM data_records`).Scan(&uniqueSources)
	if err != nil {
		return nil, fmt.Errorf("failed to get unique sources: %w", err)
	}
	stats["unique_sources"] = uniqueSources

	// Количество уникальных ключей
	var uniqueKeys int64
	err = db.pool.QueryRow(ctx, `SELECT COUNT(DISTINCT key) FROM data_records`).Scan(&uniqueKeys)
	if err != nil {
		return nil, fmt.Errorf("failed to get unique keys: %w", err)
	}
	stats["unique_keys"] = uniqueKeys

	// Количество уникальных значений
	var uniqueValues int64
	err = db.pool.QueryRow(ctx, `SELECT COUNT(DISTINCT value) FROM data_records`).Scan(&uniqueValues)
	if err != nil {
		return nil, fmt.Errorf("failed to get unique values: %w", err)
	}
	stats["unique_values"] = uniqueValues

	// Количество записей за последние 24 часа
	var recentRecords int64
	err = db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM data_records WHERE date_first >= NOW() - INTERVAL '24 hours'`).Scan(&recentRecords)
	if err != nil {
		return nil, fmt.Errorf("failed to get recent records: %w", err)
	}
	stats["recent_records"] = recentRecords

	return stats, nil
}

// GetDatabaseSize - получает размер базы данных
func (db *PostgreSQL) GetDatabaseSize() (int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var size int64
	err := db.pool.QueryRow(ctx, `SELECT pg_database_size(current_database())`).Scan(&size)
	if err != nil {
		return 0, fmt.Errorf("failed to get database size: %w", err)
	}

	return size, nil
}

// ExecuteQuery - выполняет произвольный SQL-запрос
func (db *PostgreSQL) ExecuteQuery(query string, args ...interface{}) (sql.Result, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	return db.pool.Exec(ctx, query, args...)
}

// Query - выполняет произвольный SQL-запрос и возвращает результаты
func (db *PostgreSQL) Query(query string, args ...interface{}) (*pgxpool.Rows, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	return db.pool.Query(ctx, query, args...)
}

// Begin - начинает транзакцию
func (db *PostgreSQL) Begin() (pgxpool.Tx, error) {
	return db.pool.Begin(context.Background())
}

// GetPool - получает пул соединений
func (db *PostgreSQL) GetPool() *pgxpool.Pool {
	return db.pool
}

// DataRecord - запись данных
type DataRecord struct {
	ID              int64
	SourceName      string
	Key             string
	Value           string
	FirstAppearance time.Time
	LastAppearance  time.Time
}