country-geo-locations/pkg/geoloc/database.go

66 lines
1.1 KiB
Go

package geoloc
import (
"fmt"
"github.com/hashicorp/go-memdb"
)
var database *memdb.MemDB
func CreateDatabase(ipinfos []IPInfo) (err error) {
schema := &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"},
},
},
},
},
}
database, err = memdb.NewMemDB(schema)
if err != nil {
return
}
txn := database.Txn(true)
defer txn.Abort()
for _, ipinfo := range ipinfos {
if err = txn.Insert("ipinfo", ipinfo); err != nil {
return
}
}
txn.Commit()
return
}
func SearchIPNet(ipnetnum uint) (*IPInfo, error) {
txn := database.Txn(false)
defer txn.Abort()
var ipinfo IPInfo
for {
raw, err := txn.First("ipinfo", "id", ipnetnum)
if err != nil {
return nil, err
}
if raw != nil {
ipinfo = raw.(IPInfo)
break
}
if ipnetnum == 0 {
return nil, fmt.Errorf("SearchIPNet: IP net not found")
}
ipnetnum = ipnetnum - 256
}
return &ipinfo, nil
}