country-geo-locations/pkg/geoloc/database.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
}