diff --git a/internal/database/actions.go b/internal/database/actions.go index 6f7dc77..6b9c93f 100644 --- a/internal/database/actions.go +++ b/internal/database/actions.go @@ -125,3 +125,49 @@ func CreateNetwork(sess sqlbuilder.Database, location *utils.LocationInfo) (*Net network.ID = newID.(int64) return &network, nil } + +// CreateOrUpdateURL will create a new URL entry to the urls table if it doesn't +// exists, otherwise it will update the category code of the one already in +// there. +func CreateOrUpdateURL(sess sqlbuilder.Database, url string, categoryCode string, countryCode string) (int64, error) { + var urlID int64 + + res, err := sess.Update("urls").Set( + "url", url, + "category_code", categoryCode, + "country_code", countryCode, + ).Where("url = ? AND country_code = ?", url, countryCode).Exec() + + if err != nil { + log.Error("Failed to write to the URL table") + return 0, err + } + affected, err := res.RowsAffected() + + if err != nil { + log.Error("Failed to get affected row count") + return 0, err + } + if affected == 0 { + newID, err := sess.Collection("urls").Insert( + URL{ + URL: sql.NullString{String: url, Valid: true}, + CategoryCode: sql.NullString{String: categoryCode, Valid: true}, + CountryCode: sql.NullString{String: countryCode, Valid: true}, + }) + if err != nil { + log.Error("Failed to insert into the URLs table") + return 0, err + } + urlID = newID.(int64) + } else { + lastID, err := res.LastInsertId() + if err != nil { + log.Error("failed to get URL ID") + return 0, err + } + urlID = lastID + } + + return urlID, nil +} diff --git a/nettests/websites/web_connectivity.go b/nettests/websites/web_connectivity.go index b4db6d4..32a7bf0 100644 --- a/nettests/websites/web_connectivity.go +++ b/nettests/websites/web_connectivity.go @@ -53,41 +53,10 @@ func lookupURLs(ctl *nettests.Controller) ([]string, map[int64]int64, error) { } for idx, url := range parsed.Results { - var urlID int64 - - res, err := ctl.Ctx.DB.Update("urls").Set( - "url", url.URL, - "category_code", url.CategoryCode, - "country_code", url.CountryCode, - ).Where("url = ? AND country_code = ?", url.URL, url.CountryCode).Exec() - + urlID, err := database.CreateOrUpdateURL(ctl.Ctx.DB, url.URL, url.CategoryCode, url.CountryCode) if err != nil { - log.Error("Failed to write to the URL table") - } else { - affected, err := res.RowsAffected() - - if err != nil { - log.Error("Failed to get affected row count") - } else if affected == 0 { - newID, err := ctl.Ctx.DB.Collection("urls").Insert( - database.URL{ - URL: url.URL, - CategoryCode: url.CategoryCode, - CountryCode: url.CountryCode, - }) - if err != nil { - log.Error("Failed to insert into the URLs table") - } - urlID = newID.(int64) - } else { - lastID, err := res.LastInsertId() - if err != nil { - log.Error("failed to get URL ID") - } - urlID = lastID - } + log.Error("failed to add to the URL table") } - urlIDMap[int64(idx)] = urlID urls = append(urls, url.URL) }