124 lines
2.3 KiB
Go
124 lines
2.3 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/uptrace/bun"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type Cache interface {
|
|
Get(id uint64) []string
|
|
Append(id uint64, ip string)
|
|
}
|
|
|
|
type ConnLog struct {
|
|
bun.BaseModel `bun:"table:conn_log"`
|
|
UserID uint64 `bun:"user_id"`
|
|
IP string `bun:"ip_addr"`
|
|
TS time.Time `bun:"ts"`
|
|
}
|
|
|
|
func NewConnLogs(ctx context.Context, db *bun.DB, cache Cache, logger *zap.Logger, updateInterval time.Duration) (*ConnLogs, error) {
|
|
connLogs := &ConnLogs{db: db, cache: cache}
|
|
|
|
logger.Info("filling initial cache")
|
|
err := connLogs.fillCache(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
logger.Info("initial cache filled")
|
|
|
|
go func() {
|
|
ticker := time.NewTicker(updateInterval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
var err error
|
|
|
|
err = connLogs.fillCache(ctx)
|
|
if err != nil {
|
|
logger.Error("update cache", zap.Error(err))
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
return connLogs, nil
|
|
}
|
|
|
|
type ConnLogs struct {
|
|
db *bun.DB
|
|
cache Cache
|
|
lastTS time.Time
|
|
}
|
|
|
|
func (l *ConnLogs) fillCache(ctx context.Context) error {
|
|
var entity []ConnLog
|
|
|
|
query := l.db.NewSelect().Model(&entity).
|
|
Order("ts").
|
|
Group("user_id", "ip_addr").
|
|
Column("user_id", "ip_addr").
|
|
ColumnExpr(`max("ts") as ts`)
|
|
|
|
if !l.lastTS.IsZero() {
|
|
query.Where(`"ts" > ? `, l.lastTS)
|
|
}
|
|
|
|
if err := query.Scan(ctx); err != nil {
|
|
return fmt.Errorf("select: %w", err)
|
|
}
|
|
|
|
for i := range entity {
|
|
l.cache.Append(entity[i].UserID, entity[i].IP)
|
|
}
|
|
|
|
if len(entity) > 0 {
|
|
l.lastTS = entity[len(entity)-1].TS
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (l *ConnLogs) Get(_ context.Context, first, second uint64) (bool, error) {
|
|
ips1 := l.cache.Get(first)
|
|
ips2 := l.cache.Get(second)
|
|
|
|
if len(ips1) == 0 || len(ips2) == 0 {
|
|
return false, nil
|
|
}
|
|
|
|
for i := range ips1 {
|
|
for j := range ips2 {
|
|
if ips1[i] == ips2[j] {
|
|
return true, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
return false, nil
|
|
}
|
|
|
|
func (l *ConnLogs) List(ctx context.Context) (string, error) {
|
|
var entity []ConnLog
|
|
|
|
if err := l.db.NewSelect().Model(&entity).Scan(ctx); err != nil {
|
|
return "", fmt.Errorf("select: %w", err)
|
|
}
|
|
|
|
marshal, err := json.Marshal(entity)
|
|
if err != nil {
|
|
return "", fmt.Errorf("marshal: %w", err)
|
|
}
|
|
|
|
return string(marshal), nil
|
|
}
|