Improve logging

Author: paul-norman

helpers.go

@@ -22,19 +22,19 @@ import (
 func decompressFile(filePath string, compression string) error {
 	reader, err := os.Open(filePath)
 	if err != nil {
-	   panic(err)
+		panic(err)
 	}
 	defer reader.Close()
 
 	uncompressedStream, err := gzip.NewReader(reader)
 	if err != nil {
-	   panic(err)
+		panic(err)
 	}
 
 	decompressedFilePath := strings.Replace(filePath, compression, "", -1)
 	decompressedFile, err := os.Create(decompressedFilePath)
 	if err != nil {
-	   panic(err)
+		panic(err)
 	}
 	defer decompressedFile.Close()
 
@@ -209,6 +209,20 @@ func getIpVersion(ipString string) int {
 	return ipVersion
 }
 
+func getLogFrequency() int {
+	loadLogFrequency := os.Getenv("LOAD_LOG_FREQ")
+	if len(loadLogFrequency) > 0 {
+		loadLogFrequencyInt, err := strconv.Atoi(loadLogFrequency)
+		if err != nil {
+			panic(err)
+		}
+
+		return loadLogFrequencyInt
+	}
+
+	return 1000
+}
+
 func hasASNDatabase() bool {
 	return len(os.Getenv("ASN")) > 0
 }

load.go

@@ -53,17 +53,9 @@ func loadCities(dataToLoad DataToLoad) {
 	defer csvFile.Close()
 	csvFileReader := csv.NewReader(csvFile)
 
-	version := dbQueryMaxVersion("ip_city", dataToLoad.Version) + 1
-	cities	:= []IpCity{}
-	count	:= 0;
-	lastLog	:= 0;
-	logFS	:= os.Getenv("LOAD_LOG_FREQ")
-	logFreq, err	:= strconv.Atoi(logFS)
-	if err != nil {
-		logFreq	= 100;
-	}
-
-
+	version 	:= dbQueryMaxVersion("ip_city", dataToLoad.Version) + 1
+	cities		:= []IpCity{}
+	numSaved	:= 0;
 
 	fmt.Println("rebuilding: ip_city ipv", dataToLoad.Version)
 	fmt.Print("\033[s") // Save the cursor position
@@ -78,22 +70,15 @@ func loadCities(dataToLoad DataToLoad) {
 		cities = append(cities, IpCity{ record[0], record[1], record[2], record[3], record[4], record[5], record[6], lat, lon, record[9], dataToLoad.Version, version })
 
 		if len(cities) == 100 {
-			count += len(cities);
-			if (count >= lastLog + logFreq) {
-				fmt.Print("\033[u\033[K") // Restore the cursor position and clear the line
-				fmt.Printf("Saved: %d entries\n", count)
-				lastLog = count
-			}
 			dbSaveCities(cities)
+			logEntriesConditionally(&numSaved, &cities)
 			cities = []IpCity{}
 		}
 	}
 
 	if len(cities) > 0 {
-		count += len(cities);
-		fmt.Print("\033[u\033[K") // Restore the cursor position and clear the line
-		fmt.Printf("Saved: %d entries\n", count)
 		dbSaveCities(cities)
+		logEntries(&cities)
 	}
 
 	dbDropOld("ip_city", dataToLoad.Version, version)
@@ -107,15 +92,9 @@ func loadASNs(dataToLoad DataToLoad) {
 	defer csvFile.Close()
 	csvFileReader := csv.NewReader(csvFile)
 
-	version := dbQueryMaxVersion("ip_asn", dataToLoad.Version) + 1
-	ASNs	:= []IpASN{}
-	count	:= 0;
-	lastLog	:= 0;
-	logFS	:= os.Getenv("LOAD_LOG_FREQ")
-	logFreq, err	:= strconv.Atoi(logFS)
-	if err != nil {
-		logFreq	= 100;
-	}
+	version 	:= dbQueryMaxVersion("ip_asn", dataToLoad.Version) + 1
+	ASNs		:= []IpASN{}
+	numSaved	:= 0;
 
 	fmt.Println("rebuilding: ip_asn ipv", dataToLoad.Version)
 	fmt.Print("\033[s") // Save the cursor position
@@ -128,22 +107,15 @@ func loadASNs(dataToLoad DataToLoad) {
 		ASNs = append(ASNs, IpASN{ record[0], record[1], asn, record[3], dataToLoad.Version, version })
 
 		if len(ASNs) == 100 {
-			count += len(ASNs);
-			if (count >= lastLog + logFreq) {
-				fmt.Print("\033[u\033[K") // Restore the cursor position and clear the line
-				fmt.Printf("Saved: %d entries\n", count)
-				lastLog = count
-			}
 			dbSaveASNs(ASNs)
+			logEntriesConditionally(&numSaved, &ASNs)
 			ASNs = []IpASN{}
 		}
 	}
 
 	if len(ASNs) > 0 {
-		count += len(ASNs);
-		fmt.Print("\033[u\033[K") // Restore the cursor position and clear the line
-		fmt.Printf("Saved: %d entries\n", count)
 		dbSaveASNs(ASNs)
+		logEntries(&ASNs)
 	}
 
 	dbDropOld("ip_asn", dataToLoad.Version, version)
@@ -159,13 +131,7 @@ func loadCountries(dataToLoad DataToLoad) {
 
 	version		:= dbQueryMaxVersion("ip_country", dataToLoad.Version) + 1
 	countries	:= []IpCountry{}
-	count		:= 0;
-	lastLog	:= 0;
-	logFS	:= os.Getenv("LOAD_LOG_FREQ")
-	logFreq, err	:= strconv.Atoi(logFS)
-	if err != nil {
-		logFreq	= 100;
-	}
+	numSaved	:= 0
 
 	fmt.Println("rebuilding: ip_country ipv", dataToLoad.Version)
 	fmt.Print("\033[s") // Save the cursor position
@@ -177,27 +143,36 @@ func loadCountries(dataToLoad DataToLoad) {
 		countries = append(countries, IpCountry{ record[0], record[1], record[2], dataToLoad.Version, version })
 
 		if len(countries) == 100 {
-			count += len(countries);
-			if (count >= lastLog + logFreq) {
-				fmt.Print("\033[u\033[K") // Restore the cursor position and clear the line
-				fmt.Printf("Saved: %d entries\n", count)
-				lastLog = count
-			}
 			dbSaveCountries(countries)
+			logEntriesConditionally(&numSaved, &countries)
 			countries = []IpCountry{}
 		}
 	}
 
 	if len(countries) > 0 {
-		count += len(countries);
-		fmt.Print("\033[u\033[K") // Restore the cursor position and clear the line
-		fmt.Printf("Saved: %d entries\n", count)
 		dbSaveCountries(countries)
+		logEntries(&countries)
 	}
 
 	dbDropOld("ip_country", dataToLoad.Version, version)
 }
 
 func loadDbStructure() {
 	dbFile()
+}
+
+func logEntries[T any](list *[]T) {
+	fmt.Print("\033[u\033[K") // Restore the cursor position and clear the line
+	fmt.Printf("Saved: %d entries\n", len(*list))
+}
+
+// Don't spam the logs, only display every message after `logFreq` records
+func logEntriesConditionally[T any](numSaved *int, list *[]T) {
+	logFreq := getLogFrequency()
+
+	*numSaved += len(*list)
+
+	if len(*list) >= *numSaved + logFreq {
+		logEntries(list)
+	}
 }