Make geo database usage optional (#39)

models/geo.go

@@ -1,13 +1,13 @@
 package models
 
 import (
+	"fmt"
 	"log"
 	"net"
 
 	"github.com/oschwald/maxminddb-golang"
 )
 
-// GeoRecord is the model for City database
 type GeoRecord struct {
 	Country struct {
 		ISOCode string            `maxminddb:"iso_code"`
@@ -26,52 +26,107 @@ type GeoRecord struct {
 	} `maxminddb:"postal"`
 }
 
-// ASNRecord is the model for ASN database
 type ASNRecord struct {
 	AutonomousSystemNumber       uint   `maxminddb:"autonomous_system_number"`
 	AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
 }
 
-type geodb struct {
-	city *maxminddb.Reader
-	asn  *maxminddb.Reader
+type GeoDB struct {
+	cityPath string
+	asnPath  string
+	City     *maxminddb.Reader
+	ASN      *maxminddb.Reader
 }
 
-var db geodb
+func Setup(cityPath string, asnPath string) (*GeoDB, error) {
+	city, asn, err := openDatabases(cityPath, asnPath)
+	if err != nil {
+		return nil, err
+	}
 
-func openMMDB(path string) *maxminddb.Reader {
+	return &GeoDB{
+		cityPath: cityPath,
+		asnPath:  asnPath,
+		City:     city,
+		ASN:      asn,
+	}, nil
+}
+
+func (db *GeoDB) CloseDBs() error {
+	var errs []error
+
+	if db.City != nil {
+		if err := db.City.Close(); err != nil {
+			errs = append(errs, fmt.Errorf("closing city db: %w", err))
+		}
+	}
+
+	if db.ASN != nil {
+		if err := db.ASN.Close(); err != nil {
+			errs = append(errs, fmt.Errorf("closing ASN db: %w", err))
+		}
+	}
+
+	if len(errs) > 0 {
+		return fmt.Errorf("errors closing databases: %s", errs)
+	}
+	return nil
+}
+
+func (db *GeoDB) Reload() error {
+	if err := db.CloseDBs(); err != nil {
+		return fmt.Errorf("closing existing connections: %w", err)
+	}
+
+	city, asn, err := openDatabases(db.cityPath, db.asnPath)
+	if err != nil {
+		return fmt.Errorf("opening new connections: %w", err)
+	}
+
+	db.City = city
+	db.ASN = asn
+	return nil
+}
+
+func (db *GeoDB) LookupCity(ip net.IP) (*GeoRecord, error) {
+	record := &GeoRecord{}
+	err := db.City.Lookup(ip, record)
+	if err != nil {
+		return nil, err
+	}
+
+	return record, nil
+}
+
+func (db *GeoDB) LookupASN(ip net.IP) (*ASNRecord, error) {
+	record := &ASNRecord{}
+	err := db.ASN.Lookup(ip, record)
+	if err != nil {
+		return nil, err
+	}
+	return record, nil
+}
+
+func openDatabases(cityPath, asnPath string) (*maxminddb.Reader, *maxminddb.Reader, error) {
+	city, err := openMMDB(cityPath)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	asn, err := openMMDB(asnPath)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return city, asn, nil
+}
+
+func openMMDB(path string) (*maxminddb.Reader, error) {
 	db, err := maxminddb.Open(path)
 	if err != nil {
-		log.Fatal(err)
+		return nil, err
 	}
 	log.Printf("Database %s has been loaded\n", path)
 
-	return db
-}
-
-// Setup opens all Geolite2 databases
-func Setup(cityPath string, asnPath string) {
-	db.city = openMMDB(cityPath)
-	db.asn = openMMDB(asnPath)
-}
-
-// CloseDBs unmaps from memory and frees resources to the filesystem
-func CloseDBs() {
-	log.Printf("Closing dbs...")
-	if err := db.city.Close(); err != nil {
-		log.Printf("Error closing city db: %s", err)
-	}
-	if err := db.asn.Close(); err != nil {
-		log.Printf("Error closing ASN db: %s", err)
-	}
-}
-
-// LookUp an IP and get city data
-func (record *GeoRecord) LookUp(ip net.IP) error {
-	return db.city.Lookup(ip, record)
-}
-
-// LookUp an IP and get ASN data
-func (record *ASNRecord) LookUp(ip net.IP) error {
-	return db.asn.Lookup(ip, record)
+	return db, nil
 }

models/geo_test.go

@@ -1,10 +1,12 @@
 package models
 
 import (
+	"fmt"
 	"net"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestModels(t *testing.T) {
@@ -59,19 +61,21 @@ func TestModels(t *testing.T) {
 		AutonomousSystemOrganization: "IP-Only",
 	}
 
-	Setup("../test/GeoIP2-City-Test.mmdb", "../test/GeoLite2-ASN-Test.mmdb")
-	defer CloseDBs()
+	db, err := Setup("../test/GeoIP2-City-Test.mmdb", "../test/GeoLite2-ASN-Test.mmdb")
+	require.NoError(t, err, fmt.Sprintf("Error setting up db: %s", err))
+	defer db.CloseDBs()
+	assert.NotNil(t, db.ASN)
+	assert.NotNil(t, db.City)
 
-	assert.NotNil(t, db.asn)
-	assert.NotNil(t, db.city)
-
-	cityRecord := &GeoRecord{}
-	assert.Nil(t, cityRecord.LookUp(net.ParseIP("81.2.69.192")))
+	cityRecord, err := db.LookupCity(net.ParseIP("81.2.69.192"))
+	require.NoError(t, err, fmt.Sprintf("Error looking up city: %s", err))
 	assert.Equal(t, expectedCity, cityRecord)
-	assert.Error(t, cityRecord.LookUp(net.ParseIP("error")))
+	_, err = db.LookupCity(net.ParseIP("error"))
+	assert.Error(t, err)
 
-	asnRecord := &ASNRecord{}
-	assert.Nil(t, asnRecord.LookUp(net.ParseIP("82.99.17.64")))
+	asnRecord, err := db.LookupASN(net.ParseIP("82.99.17.64"))
+	require.NoError(t, err, fmt.Sprintf("Error looking up asn: %s", err))
 	assert.Equal(t, expectedASN, asnRecord)
-	assert.Error(t, asnRecord.LookUp(net.ParseIP("error")))
+	_, err = db.LookupASN(net.ParseIP("error"))
+	assert.Error(t, err)
 }