From c7e12ba3b53f8e9346c2e179adca4d3e9a303845 Mon Sep 17 00:00:00 2001
From: jabra98 <jabra98.gpg@gmail.com>
Date: Sat, 1 Jan 2022 20:33:58 +0100
Subject: [PATCH] fix: consider all Machine/UserAgent entries

---
 models/compat/wakatime/v1/machine.go    |  3 +-
 models/compat/wakatime/v1/user_agent.go |  3 +-
 services/imports/wakatime.go            | 89 +++++++++++++++----------
 3 files changed, 59 insertions(+), 36 deletions(-)

diff --git a/models/compat/wakatime/v1/machine.go b/models/compat/wakatime/v1/machine.go
index bfdc68a..c46e1e2 100644
--- a/models/compat/wakatime/v1/machine.go
+++ b/models/compat/wakatime/v1/machine.go
@@ -3,7 +3,8 @@ package v1
 // https://wakatime.com/api/v1/users/current/machine_names
 
 type MachineViewModel struct {
-	Data []*MachineEntry `json:"data"`
+	Data       []*MachineEntry      `json:"data"`
+	TotalPages int                  `json:"total_pages"`
 }
 
 type MachineEntry struct {
diff --git a/models/compat/wakatime/v1/user_agent.go b/models/compat/wakatime/v1/user_agent.go
index af427bd..4fc127d 100644
--- a/models/compat/wakatime/v1/user_agent.go
+++ b/models/compat/wakatime/v1/user_agent.go
@@ -1,7 +1,8 @@
 package v1
 
 type UserAgentsViewModel struct {
-	Data []*UserAgentEntry `json:"data"`
+	Data       []*UserAgentEntry    `json:"data"`
+	TotalPages int                  `json:"total_pages"`
 }
 
 type UserAgentEntry struct {
diff --git a/services/imports/wakatime.go b/services/imports/wakatime.go
index 61fe731..8c8f3ce 100644
--- a/services/imports/wakatime.go
+++ b/services/imports/wakatime.go
@@ -15,6 +15,7 @@ import (
 	"golang.org/x/sync/semaphore"
 	"net/http"
 	"time"
+	"strconv"
 )
 
 const OriginWakatime = "wakatime"
@@ -174,24 +175,34 @@ func (w *WakatimeHeartbeatImporter) fetchRange() (time.Time, time.Time, error) {
 func (w *WakatimeHeartbeatImporter) fetchUserAgents() (map[string]*wakatime.UserAgentEntry, error) {
 	httpClient := &http.Client{Timeout: 10 * time.Second}
 
-	req, err := http.NewRequest(http.MethodGet, config.WakatimeApiUrl+config.WakatimeApiUserAgentsUrl, nil)
-	if err != nil {
-		return nil, err
-	}
-
-	res, err := httpClient.Do(w.withHeaders(req))
-	if err != nil {
-		return nil, err
-	}
-
-	var userAgentsData wakatime.UserAgentsViewModel
-	if err := json.NewDecoder(res.Body).Decode(&userAgentsData); err != nil {
-		return nil, err
-	}
-
 	userAgents := make(map[string]*wakatime.UserAgentEntry)
-	for _, ua := range userAgentsData.Data {
-		userAgents[ua.Id] = ua
+
+	done := false
+	for i:=1; !done; i++ {
+		queryStr := "?page="+strconv.Itoa(i)
+
+		req, err := http.NewRequest(http.MethodGet, config.WakatimeApiUrl+
+						config.WakatimeApiUserAgentsUrl+queryStr, nil)
+		if err != nil {
+			return nil, err
+		}
+
+		res, err := httpClient.Do(w.withHeaders(req))
+		if err != nil {
+			return nil, err
+		}
+
+		var userAgentsData wakatime.UserAgentsViewModel
+		if err := json.NewDecoder(res.Body).Decode(&userAgentsData); err != nil {
+			return nil, err
+		}
+		if i == userAgentsData.TotalPages {
+			done = true
+		}
+
+		for _, ua := range userAgentsData.Data {
+			userAgents[ua.Id] = ua
+		}
 	}
 
 	return userAgents, nil
@@ -202,24 +213,34 @@ func (w *WakatimeHeartbeatImporter) fetchUserAgents() (map[string]*wakatime.User
 func (w *WakatimeHeartbeatImporter) fetchMachineNames() (map[string]*wakatime.MachineEntry, error) {
 	httpClient := &http.Client{Timeout: 10 * time.Second}
 
-	req, err := http.NewRequest(http.MethodGet, config.WakatimeApiUrl+config.WakatimeApiMachineNamesUrl, nil)
-	if err != nil {
-		return nil, err
-	}
-
-	res, err := httpClient.Do(w.withHeaders(req))
-	if err != nil {
-		return nil, err
-	}
-
-	var machineData wakatime.MachineViewModel
-	if err := json.NewDecoder(res.Body).Decode(&machineData); err != nil {
-		return nil, err
-	}
-
 	machines := make(map[string]*wakatime.MachineEntry)
-	for _, ma := range machineData.Data {
-		machines[ma.Id] = ma
+
+	done := false
+	for i:=1; !done; i++ {
+		queryStr := "?page="+strconv.Itoa(i)
+
+		req, err := http.NewRequest(http.MethodGet, config.WakatimeApiUrl+
+						config.WakatimeApiMachineNamesUrl+queryStr, nil)
+		if err != nil {
+			return nil, err
+		}
+
+		res, err := httpClient.Do(w.withHeaders(req))
+		if err != nil {
+			return nil, err
+		}
+
+		var machineData wakatime.MachineViewModel
+		if err := json.NewDecoder(res.Body).Decode(&machineData); err != nil {
+				return nil, err
+		}
+		if i == machineData.TotalPages {
+			done = true
+		}
+
+		for _, ma := range machineData.Data {
+			machines[ma.Id] = ma
+		}
 	}
 
 	return machines, nil