Ensure partial ASN batch is written

internal/metrics/crawler.go

@@ -199,10 +199,12 @@ func (s *CrawlerServiceMetrics) StartIPMapping(limit uint) {
 func (s *CrawlerServiceMetrics) GetIPsAfterTimestamp(ips *rpc.GetIPsAfterTimestampResponse) {
 	// If we don't have either maxmind DB, bail now
 	if s.maxMindCountryDB == nil && s.maxMindASNDB == nil {
+		log.Debug("Missing both ASN and Country maxmind DBs")
 		return
 	}
 
 	if ips == nil {
+		log.Debug("IPs after timestamp are nil")
 		return
 	}
 
@@ -250,6 +252,7 @@ func (s *CrawlerServiceMetrics) ProcessIPCountryMapping(ips *rpc.GetIPsAfterTime
 func (s *CrawlerServiceMetrics) ProcessIPASNMapping(ips *rpc.GetIPsAfterTimestampResponse) {
 	// Don't process if we can't store
 	if s.metrics.mysqlClient == nil {
+		log.Debug("MySQL client is nil")
 		return
 	}
 	if s.network == nil {
@@ -264,10 +267,13 @@ func (s *CrawlerServiceMetrics) ProcessIPASNMapping(ips *rpc.GetIPsAfterTimestam
 	}
 	asnCounts := map[uint32]*countStruct{}
 
+	log.Debugf("Have %d IPs\n", len(ips.IPs.MustGet()))
+
 	if ipresult, hasIPResult := ips.IPs.Get(); hasIPResult {
 		for _, ip := range ipresult {
 			asn, err := s.GetASNForIP(ip)
 			if err != nil {
+				log.Debugf("Unable to get ASN for IP %s\n", ip)
 				continue
 			}
 
@@ -303,7 +309,7 @@ func (s *CrawlerServiceMetrics) ProcessIPASNMapping(ips *rpc.GetIPsAfterTimestam
 
 		// Execute the batch insert when reaching the batch size or the end of the slice
 		if (i+1)%batchSize == 0 || i+1 == uint32(len(asnCounts)) {
-			_, err := s.metrics.mysqlClient.Exec(
+			_, err = s.metrics.mysqlClient.Exec(
 				fmt.Sprintf("INSERT INTO asn(asn, organization, count, network) VALUES %s", strings.Join(valueStrings, ",")),
 				valueArgs...)
 
@@ -316,6 +322,17 @@ func (s *CrawlerServiceMetrics) ProcessIPASNMapping(ips *rpc.GetIPsAfterTimestam
 			valueArgs = []interface{}{}
 		}
 	}
+
+	// If there are any records left that did not meet the batch size threshold, insert them now
+	if len(valueStrings) > 0 {
+		_, err := s.metrics.mysqlClient.Exec(
+			fmt.Sprintf("INSERT INTO asn(asn, organization, count, network) VALUES %s", strings.Join(valueStrings, ",")),
+			valueArgs...)
+
+		if err != nil {
+			log.Errorf("error inserting remaining ASN records to mysql: %s\n", err.Error())
+		}
+	}
 }
 
 // CountryRecord record of a country from maxmind

internal/metrics/database.go

@@ -10,7 +10,7 @@ func (m *Metrics) initTables() error {
 		"  `organization` VARCHAR(255) NOT NULL," +
 		"  `count` int unsigned NOT NULL," +
 		"  `network` VARCHAR(255) NOT NULL," +
-		"UNIQUE KEY `asn-unique` (`asn`)" +
+		"UNIQUE KEY `asn-network-unique` (`asn`, `network`)" +
 		") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci;"
 
 	result, err := m.mysqlClient.Query(query)

internal/metrics/metrics.go

@@ -109,11 +109,11 @@ func NewMetrics(port uint16, logLevel log.Level) (*Metrics, error) {
 
 	err = metrics.createDBClient()
 	if err != nil {
-		log.Debugf("ERROR creating MySQL Client. Will not store any metrics to MySQL")
+		log.Debugf("Error creating MySQL Client. Will not store any metrics to MySQL. %s\n", err.Error())
 	}
 	err = metrics.initTables()
 	if err != nil {
-		log.Debugf("ERROR ensuring tables exist in MySQL. Will not process or store any MySQL only metrics")
+		log.Errorf("Error ensuring tables exist in MySQL. Will not process or store any MySQL only metrics: %s\n", err.Error())
 		metrics.mysqlClient = nil
 	}
 

internal/metrics/promhandler.go

@@ -1,4 +1,5 @@
 package metrics
+
 // Everything in here is to support making a new prometheus registry and having the http endpoint pick it up seamlessly
 
 import (