Make geo database usage optional (#39)

service/geo.go

@@ -1,37 +1,73 @@
 package service
 
 import (
+	"context"
 	"log"
 	"net"
+	"sync"
 
 	"github.com/dcarrillo/whatismyip/models"
 )
 
-// Geo defines a base type for lookups
 type Geo struct {
-	IP net.IP
+	ctx    context.Context
+	cancel context.CancelFunc
+	db     *models.GeoDB
+	mu     sync.RWMutex
 }
 
-// LookUpCity queries the database for city data related to the given IP
-func (g *Geo) LookUpCity() *models.GeoRecord {
-	record := &models.GeoRecord{}
-	err := record.LookUp(g.IP)
+func NewGeo(ctx context.Context, cityPath string, asnPath string) (*Geo, error) {
+	ctx, cancel := context.WithCancel(ctx)
+
+	db, err := models.Setup(cityPath, asnPath)
 	if err != nil {
-		log.Println(err)
+		cancel()
+		return nil, err
+	}
+
+	geo := &Geo{
+		ctx:    ctx,
+		cancel: cancel,
+		db:     db,
+	}
+
+	return geo, nil
+}
+
+func (g *Geo) LookUpCity(ip net.IP) *models.GeoRecord {
+	record, err := g.db.LookupCity(ip)
+	if err != nil {
+		log.Print(err)
 		return nil
 	}
 
 	return record
 }
 
-// LookUpASN queries the database for ASN data related to the given IP
-func (g *Geo) LookUpASN() *models.ASNRecord {
-	record := &models.ASNRecord{}
-	err := record.LookUp(g.IP)
+func (g *Geo) LookUpASN(ip net.IP) *models.ASNRecord {
+	record, err := g.db.LookupASN(ip)
 	if err != nil {
-		log.Println(err)
+		log.Print(err)
 		return nil
 	}
 
 	return record
 }
+
+func (g *Geo) Shutdown() {
+	g.cancel()
+	g.db.CloseDBs()
+}
+
+func (g *Geo) Reload() {
+	if err := g.ctx.Err(); err != nil {
+		log.Printf("Skipping reload, service is shutting down: %v", err)
+		return
+	}
+
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	g.db.Reload()
+	log.Print("Geo database reloaded")
+}

service/geo_test.go

@@ -1,36 +1,33 @@
 package service
 
 import (
+	"context"
 	"net"
 	"os"
 	"testing"
 
-	"github.com/dcarrillo/whatismyip/models"
 	"github.com/stretchr/testify/assert"
 )
 
+var geoSvc *Geo
+
 func TestMain(m *testing.M) {
-	models.Setup("../test/GeoIP2-City-Test.mmdb", "../test/GeoLite2-ASN-Test.mmdb")
-	defer models.CloseDBs()
+	geoSvc, _ = NewGeo(context.Background(), "../test/GeoIP2-City-Test.mmdb", "../test/GeoLite2-ASN-Test.mmdb")
 	os.Exit(m.Run())
 }
 
 func TestCityLookup(t *testing.T) {
-	ip := Geo{IP: net.ParseIP("error")}
-	c := ip.LookUpCity()
+	c := geoSvc.LookUpCity(net.ParseIP("error"))
 	assert.Nil(t, c)
 
-	ip = Geo{IP: net.ParseIP("1.1.1.1")}
-	c = ip.LookUpCity()
+	c = geoSvc.LookUpCity(net.ParseIP("1.1.1.1"))
 	assert.NotNil(t, c)
 }
 
 func TestASNLookup(t *testing.T) {
-	ip := Geo{IP: net.ParseIP("error")}
-	a := ip.LookUpASN()
+	a := geoSvc.LookUpASN(net.ParseIP("error"))
 	assert.Nil(t, a)
 
-	ip = Geo{IP: net.ParseIP("1.1.1.1")}
-	a = ip.LookUpASN()
+	a = geoSvc.LookUpASN(net.ParseIP("1.1.1.1"))
 	assert.NotNil(t, a)
 }