country-geo-locations/internal/database/database.go

116 lines
2 KiB
Go

package database
import (
"errors"
"os"
"time"
"git.ar21.de/yolokube/country-geo-locations/internal/cmd"
"git.ar21.de/yolokube/country-geo-locations/pkg/geoloc"
"github.com/hashicorp/go-memdb"
)
var (
ErrUnknownInterface = errors.New("unknown interface structure")
ErrIPNetNotFound = errors.New("IP net not found")
)
type Database struct {
ready bool
config *cmd.AppSettings
db *memdb.MemDB
}
func NewDatabase(config *cmd.AppSettings) (*Database, error) {
database, err := memdb.NewMemDB(
&memdb.DBSchema{
Tables: map[string]*memdb.TableSchema{
"ipinfo": {
Name: "ipinfo",
Indexes: map[string]*memdb.IndexSchema{
"id": {
Name: "id",
Unique: true,
Indexer: &memdb.UintFieldIndex{Field: "IPNumFrom"},
},
},
},
},
},
)
if err != nil {
return nil, err
}
return &Database{
ready: false,
config: config,
db: database,
}, nil
}
func (d *Database) Load(ipinfos []geoloc.IPInfo) error {
txn := d.db.Txn(true)
defer txn.Abort()
for _, ipinfo := range ipinfos {
if err := txn.Insert("ipinfo", ipinfo); err != nil {
return err
}
}
txn.Commit()
d.ready = true
return nil
}
func (d *Database) IsReady() bool {
return d.ready
}
func (d *Database) SearchIPNet(ipnetnum uint) (*geoloc.IPInfo, error) {
txn := d.db.Txn(false)
defer txn.Abort()
var (
ipinfo geoloc.IPInfo
ok bool
)
for {
raw, err := txn.First("ipinfo", "id", ipnetnum)
if err != nil {
return nil, err
}
if raw != nil {
ipinfo, ok = raw.(geoloc.IPInfo)
if !ok {
return nil, ErrUnknownInterface
}
break
}
if ipnetnum == 0 {
return nil, ErrIPNetNotFound
}
ipnetnum -= geoloc.CalculationValue
}
return &ipinfo, nil
}
func (d *Database) Timestamp() (time.Time, error) {
file, err := os.Open(d.config.DataFile)
if err != nil {
return time.Time{}, err
}
stats, err := file.Stat()
if err != nil {
return time.Time{}, err
}
return stats.ModTime(), nil
}