From 714cae259442f937a02aa7011a04e81293b27e4a Mon Sep 17 00:00:00 2001 From: Steven Polley Date: Sat, 20 Jun 2020 20:24:51 -0600 Subject: [PATCH] Improved styling and comments --- countersql/database.go | 559 +---------------------------------------- main.go | 25 +- 2 files changed, 22 insertions(+), 562 deletions(-) diff --git a/countersql/database.go b/countersql/database.go index 77ce4db..1901fed 100644 --- a/countersql/database.go +++ b/countersql/database.go @@ -32,7 +32,6 @@ func (conn Connection) HasIPVisited(ipAddress string) (bool, error) { // GetUniqueVisits counts the number of entires in the visits table, representing one unique source IP address per row func (conn Connection) GetUniqueVisits() (int, error) { - rows, err := conn.DB.Query(`SELECT COUNT(*) FROM visit`) if err != nil { return 0, fmt.Errorf("SELECT query failed: %v", err) @@ -72,569 +71,15 @@ func (conn Connection) AddVisitor(ipAddress string) error { return nil } -// - -/* - -// ************************************ -// Task Database Functions -// ************************************ - -// GetTasks returns a slice of all tasks in the database. It requires -// a boolean parameter. If true, it only returns active tasks. If false -// it will return all tasks. A task is active if the time now is past -// the planned_start date. -func (conn Connection) GetTasks(activeTasksOnly bool) ([]Task, error) { - var rows *sql.Rows - var err error - if activeTasksOnly { - rows, err = conn.DB.Query(`SELECT - tasks.id, tasks.name, tasks.weight, tasks.timecost, tasks.planned_start, tasks.planned_end, tasks.actual_start, tasks.actual_end, tasks.expire_action, tasks.close_code, lists.id, lists.name - FROM tasks - INNER JOIN lists - ON tasks.list_id = lists.id - WHERE tasks.close_code IS NULL AND tasks.planned_start < ? - ORDER BY tasks.actual_start DESC, tasks.weight DESC`, time.Now().Add(-time.Hour*6)) - - } else { - rows, err = conn.DB.Query(`SELECT - tasks.id, tasks.name, tasks.weight, tasks.timecost, tasks.planned_start, tasks.planned_end, tasks.actual_start, tasks.actual_end, tasks.expire_action, tasks.close_code, lists.id, lists.name - FROM tasks - INNER JOIN lists - ON tasks.list_id = lists.id - ORDER BY tasks.actual_start DESC, tasks.weight DESC`) - } - if err != nil { - return nil, fmt.Errorf("SELECT query failed: %v", err) - } - defer rows.Close() - - if !rows.Next() { - return nil, nil - } - - var tasks []Task - for { - task := Task{} - if err := rows.Scan(&task.ID, &task.Name, &task.Weight, &task.TimeCost, &task.PlannedStart, &task.PlannedEnd, &task.ActualStart, &task.ActualEnd, &task.ExpireAction, &task.CloseCode, &task.List.ID, &task.List.Name); err != nil { - return nil, fmt.Errorf("failed to scan database row: %v", err) - } - tasks = append(tasks, task) - if !rows.Next() { - break - } - } - return tasks, nil -} - -// GetTasksByListID returns a slice of all tasks on a specific list in the database -func (conn Connection) GetTasksByListID(listID int, activeTasksOnly bool) ([]Task, error) { - activeTasksQuery := "" - if activeTasksOnly { - activeTasksQuery = " AND tasks.close_code IS NULL" - } - rows, err := conn.DB.Query(`SELECT - tasks.id, tasks.name, tasks.weight, tasks.timecost, tasks.planned_start, tasks.planned_end, tasks.actual_start, tasks.actual_end, tasks.expire_action, tasks.close_code, lists.id, lists.name - FROM tasks - INNER JOIN lists ON tasks.list_id = list.id - WHERE lists.id = ? `+activeTasksQuery, listID) - if err != nil { - return nil, fmt.Errorf("SELECT query failed: %v", err) - } - defer rows.Close() - - if !rows.Next() { - return nil, nil - } - - var tasks []Task - for { - task := Task{} - if err := rows.Scan(&task.ID, &task.Name, &task.Weight, &task.TimeCost, &task.PlannedStart, &task.PlannedEnd, &task.ActualStart, &task.ActualEnd, &task.ExpireAction, &task.CloseCode, &task.List.ID, &task.List.Name); err != nil { - return nil, fmt.Errorf("failed to scan database row: %v", err) - } - tasks = append(tasks, task) - if !rows.Next() { - break - } - } - return tasks, nil -} - -// GetTaskByID returns a task matching the ID provided -func (conn Connection) GetTaskByID(taskID int) (*Task, error) { - rows, err := conn.DB.Query(`SELECT - tasks.id, tasks.name, tasks.weight, tasks.timecost, tasks.planned_start, tasks.planned_end, tasks.actual_start, tasks.actual_end, tasks.expire_action, tasks.close_code, lists.id, lists.name - FROM tasks - INNER JOIN lists ON tasks.list_id = list.id - WHERE tasks.id = ? `, taskID) - if err != nil { - return nil, fmt.Errorf("SELECT query failed: %v", err) - } - defer rows.Close() - - if !rows.Next() { - return nil, fmt.Errorf("no results returned for taskID '%d'", taskID) - } - - task := Task{} - if err := rows.Scan(&task.ID, &task.Name, &task.Weight, &task.TimeCost, &task.PlannedStart, &task.PlannedEnd, &task.ActualStart, &task.ActualEnd, &task.ExpireAction, &task.CloseCode, &task.List.ID, &task.List.Name); err != nil { - return nil, fmt.Errorf("failed to scan database row: %v", err) - } - - return &task, nil -} - -// NewTask will insert a row into the database, given a pre-instantiated task -func (conn Connection) NewTask(task *Task) error { - _, err := conn.DB.Exec(`INSERT INTO tasks - (list_id, name, weight, timecost, planned_start, planned_end, actual_start, actual_end, expire_action, close_code) - VALUES (?, ?, ?, ?, ?, ?, NULL, NULL, ?, NULL)`, task.List.ID, task.Name, task.Weight, task.TimeCost, task.PlannedStart, task.PlannedEnd, task.ExpireAction) - if err != nil { - return fmt.Errorf("INSERT query failed: %v", err) - } - - return nil -} - -// DeleteTaskByID will remove a row from the database, given a list ID -func (conn Connection) DeleteTaskByID(taskID int) error { - _, err := conn.DB.Exec("DELETE FROM tasks WHERE tasks.id = ?", taskID) - - if err != nil { - return fmt.Errorf("Delete query failed: %v", err) - } - - return nil -} - -// StartTask will set the actual_start to now, given a task ID. This marks the task as active and started -func (conn Connection) StartTask(taskID int) error { - _, err := conn.DB.Exec("UPDATE tasks SET actual_start = now() WHERE tasks.id = ?", taskID) - - if err != nil { - return fmt.Errorf("Delete query failed: %v", err) - } - - return nil -} - -// UpdateTask will update a task row in the database, given a pre-instantiated task -func (conn Connection) UpdateTask(task *Task) error { - - var err error - _, err = conn.DB.Exec(`UPDATE tasks SET - list_id = ?, - name = ?, - weight = ?, - timecost = ?, - planned_start = ?, - planned_end = ?, - actual_start = ?, - actual_end = ?, - expire_action = ?, - close_code = ? - WHERE id = ?`, - task.List.ID, task.Name, task.Weight, task.TimeCost, task.PlannedStart, task.PlannedEnd, task.ActualStart, task.ActualEnd, task.ExpireAction, task.CloseCode, task.ID) - if err != nil { - return fmt.Errorf("UPDATE task query failed for %d: %v", task.ID, err) - } - - return nil -} - -// StopTask will stop a task by setting the close code -// closeCode can be: completed, cancelled or reset -// completed is self explanatory. -// cancelled indicates the task itself is being canelled / skipped -// reset is not an official close code, but indicates that actual_start should be cleared (timer reset) -func (conn Connection) StopTask(taskID int, closeCode string) error { - var err error - if closeCode == "reset" { - _, err = conn.DB.Exec("UPDATE tasks SET actual_start = NULL WHERE tasks.id = ?", taskID) - } else if closeCode == "cancelled" { - _, err = conn.DB.Exec("UPDATE tasks SET actual_end = now(), close_code = 'cancelled' WHERE tasks.id = ?", taskID) - } else if closeCode == "completed" { - _, err = conn.DB.Exec("UPDATE tasks SET actual_end = now(), close_code = 'completed' WHERE tasks.id = ?", taskID) - } else { - err = fmt.Errorf("invalid close code: '%s' - must be reset, cancelled or completed", closeCode) - } - - if err != nil { - return fmt.Errorf("StopTask query failed: %v", err) - } - - return nil -} - -// ************************************ -// List Database Functions -// ************************************ - -// GetLists returns a slice of all lists in the database. -// If joinLocations is true, it will grab the locations as well -func (conn Connection) GetLists(getListLocations bool) ([]List, error) { - - rows, err := conn.DB.Query(`SELECT - lists.id, lists.name - FROM lists`) - if err != nil { - return nil, fmt.Errorf("SELECT query failed: %v", err) - } - defer rows.Close() - - if !rows.Next() { - return nil, nil - } - - var lists []List - for { - list := List{} - if err := rows.Scan(&list.ID, &list.Name); err != nil { - return nil, fmt.Errorf("failed to scan database row: %v", err) - } - lists = append(lists, list) - if !rows.Next() { - break - } - } - - if getListLocations { - wg := sync.WaitGroup{} - m := sync.Mutex{} - for i := range lists { - wg.Add(1) - go func(i int) { - defer wg.Done() - rows, err := conn.DB.Query(`SELECT - listlocations.location_id, locations.name - FROM listlocations - INNER JOIN locations - ON listlocations.location_id = locations.id - WHERE list_id = ?`, lists[i].ID) - defer rows.Close() - if err != nil { - log.Printf("SELECT listlocations query failed: %v", err) - return - } - - if !rows.Next() { - return - } - - var locations []Location - for { - var location Location - if err := rows.Scan(&location.ID, &location.Name); err != nil { - log.Printf("failed to scan database row: %v", err) - continue - } - locations = append(locations, location) - if !rows.Next() { - break - } - } - m.Lock() - lists[i].Locations = locations - m.Unlock() - }(i) - } - wg.Wait() // wait - } - - return lists, nil -} - -// GetListByID will return a single list in the database with the given ID -func (conn Connection) GetListByID(listID int, getListLocations bool) (*List, error) { - rows, err := conn.DB.Query(`SELECT - lists.name - FROM lists - WHERE lists.id = ? - LIMIT 1`, listID) - if err != nil { - return nil, fmt.Errorf("SELECT query failed: %v", err) - } - defer rows.Close() - - if !rows.Next() { - return nil, fmt.Errorf("no results returned for listID '%d'", listID) - } - - list := List{} - if err := rows.Scan(&list.Name); err != nil { - return nil, fmt.Errorf("failed to scan database row: %v", err) - } - list.ID = listID - - if getListLocations { - rows, err := conn.DB.Query(`SELECT - listlocations.location_id - FROM listlocations - WHERE list_id = ?`, list.ID) - defer rows.Close() - if err != nil { - return nil, fmt.Errorf("SELECT listlocations query failed: %v", err) - } - - if !rows.Next() { - return nil, fmt.Errorf("no results returned for listID '%d'", listID) - } - - var locations []Location - for { - var location Location - if err := rows.Scan(&location.ID); err != nil { - log.Printf("failed to scan database row: %v", err) - break - } - locations = append(locations, location) - if !rows.Next() { - break - } - } - list.Locations = locations - } - - return &list, nil -} - -// NewList will insert a row into the database, given a pre-instantiated list -func (conn Connection) NewList(list *List) error { - result, err := conn.DB.Exec(`INSERT - INTO lists - (name) - VALUES (?)`, list.Name) - if err != nil { - return fmt.Errorf("INSERT list query failed: %v", err) - } - - lastInsertID, _ := result.LastInsertId() - - query := ` INSERT - INTO listlocations - (list_id, location_id) - VALUES ` - // TBD: Handle list location mapping inserts here, loop through and build query, then execute it. - values := []interface{}{} - for _, listLocation := range list.Locations { - values = append(values, lastInsertID, listLocation.ID) - // the rest of this loop is magic lifted from https://play.golang.org/p/YqNJKybpwWB - numFields := 2 - - query += `(` - for j := 0; j < numFields; j++ { - query += `?,` - } - query = query[:len(query)-1] + `),` - } - query = query[:len(query)-1] // this is also part of the above mentioned magic - _, err = conn.DB.Exec(query, values...) - if err != nil { - return fmt.Errorf("INSERT listlocations query failed: %v", err) - } - return nil -} - -// UpdateList will update a list row in the database, given a pre-instantiated list -func (conn Connection) UpdateList(list *List) error { - wg := sync.WaitGroup{} - wg.Add(2) - var err error - go func() { - defer wg.Done() - _, err = conn.DB.Exec(`UPDATE lists - (name) - VALUES (?)`, list.Name) - if err != nil { - err = fmt.Errorf("INSERT list query failed: %v", err) - return - } - }() - - // Handle list locations... Have fun - go func() { - defer wg.Done() - var rows *sql.Rows - rows, err = conn.DB.Query(`SELECT - listlocations.location_id - WHERE list_id = ?`, list.ID) - if err != nil { - err = fmt.Errorf("SELECT listlocations query failed: %v", err) - return - } - defer rows.Close() - - var locationIDsToRemove, locationIDsToAdd, dbIDs []int - - if rows.Next() { // If this is false, the list has no assigned locations in the database currently - - // We loop through databaseRows and determine locations that need to be removed - for { - var currDBID int - if err := rows.Scan(&currDBID); err != nil { - err = fmt.Errorf("failed to scan listlocations database row: %v", err) - return - } - matchFound := false - for _, currNewLocation := range list.Locations { - if currNewLocation.ID == currDBID { // If we find a match, we know this is a location that both exists in the database, and is set on the new object, so we wouldn't take action - matchFound = true - continue - } - } - - if !matchFound { // If we don't find a match, we know that this row in the database should be removed because it's no longer set on the new listObject - locationIDsToRemove = append(locationIDsToRemove, currDBID) - } - // We collect the locationIDs while looping through the database rows so we can re-use them below - dbIDs = append(dbIDs, currDBID) - - if !rows.Next() { - break - } - } - - // We loop through the newly set locations and determine any that need to be added to the database - for _, currNewLocation := range list.Locations { - matchFound := false - for _, currDBID := range dbIDs { - if currNewLocation.ID == currDBID { // If we find a match, we know the location is set on the new object and already exists in the database, so we take no action - matchFound = true - continue - } - } - - if !matchFound { // If we don't find a match, we know the location is set on the new object, but doesn't exist in the database, so we need to add it - locationIDsToAdd = append(locationIDsToAdd, currNewLocation.ID) - } - } - } else { // If there are no rows in the database then we know we just need to all all locations set on this new list object and put them in the database - for _, currNewLocation := range list.Locations { - locationIDsToAdd = append(locationIDsToAdd, currNewLocation.ID) - } - } - - // We now have populated locationIDsToAdd and locationIDsToRemove, so we just need to build and execute the queries - wg.Add(2) - go func() { - defer wg.Done() - query := ` INSERT - INTO listlocations - (list_id, location_id) - VALUES ` - // TBD: Handle list location mapping inserts here, loop through and build query, then execute it. - values := []interface{}{} - for _, locationID := range locationIDsToAdd { - values = append(values, list.ID, locationID) - // the rest of this loop is magic lifted from https://play.golang.org/p/YqNJKybpwWB - numFields := 2 - - query += `(` - for j := 0; j < numFields; j++ { - query += `?,` - } - query = query[:len(query)-1] + `),` - } - query = query[:len(query)-1] // this is also part of the above mentioned magic - _, err = conn.DB.Exec(query, values...) - if err != nil { - err = fmt.Errorf("INSERT listlocations query failed: %v", err) - return - } - }() - - go func() { - defer wg.Done() - query := ` DELETE - FROM listlocations - WHERE ` - values := []interface{}{} - for _, locationID := range locationIDsToRemove { - values = append(values, list.ID, locationID) - - query += `(list_id = ? AND location_id = ?) OR` - } - query = query[:len(query)-3] - _, err = conn.DB.Exec(query, values...) - if err != nil { - err = fmt.Errorf("DELETE listlocations query failed: %v", err) - return - } - }() - }() - wg.Wait() - if err != nil { - return fmt.Errorf("a goroutine in UpdateList failed: %v", err) - } - - return nil -} - -// DeleteListByID will remove a row from the database, given a list ID -func (conn Connection) DeleteListByID(listID int) error { - _, err := conn.DB.Exec("DELETE FROM lists WHERE lists.id = ?", listID) - - if err != nil { - return fmt.Errorf("Delete query failed: %v", err) - } - - return nil -} - -// ************************************ -// Location Database Functions -// ************************************ - -// GetLocations returns a slice of all locations in the database -func (conn Connection) GetLocations() ([]Location, error) { - rows, err := conn.DB.Query(`SELECT - locations.id, locations.name, locations.mac_address - FROM locations`) - if err != nil { - return nil, fmt.Errorf("SELECT query failed: %v", err) - } - defer rows.Close() - - if !rows.Next() { - return nil, nil - } - - var locations []Location - for { - location := Location{} - if err := rows.Scan(&location.ID, &location.Name, &location.MACAddress); err != nil { - return nil, fmt.Errorf("failed to scan database row: %v", err) - } - locations = append(locations, location) - if !rows.Next() { - break - } - } - return locations, nil -} - -// DeleteLocationByID will remove a row from the database, given a list ID -func (conn Connection) DeleteLocationByID(locationID int) error { - _, err := conn.DB.Exec(`DELETE - FROM locations - WHERE locations.id = ?`, locationID) - if err != nil { - return fmt.Errorf("Delete query failed: %v", err) - } - - return nil -} - -*/ - // Connect will open a TCP connection to the database with the given DSN configuration func (conf Configuration) Connect() (*Connection, error) { conn := &Connection{} var err error + conn.DB, err = sql.Open("mysql", conf.DSN) if err != nil { return nil, fmt.Errorf("failed to open db: %v", err) } + return conn, nil } diff --git a/main.go b/main.go index e898221..7198959 100644 --- a/main.go +++ b/main.go @@ -1,3 +1,11 @@ +/* +This application ensures the required environment variables are set +Then it uses the included countersql package to test a connection to the database +It then caches the number of unique visits in memory by querying the database +Finally, it sets up a web server with the endpoint http://localhost:8080/ +For each unique source IP address which makes a GET request to this endpoint, +the count of unique visitors is incremented by 1, both in the database and in memory cache. +*/ package main import ( @@ -12,8 +20,8 @@ import ( ) var ( - database countersql.Configuration - uniqueVisits int + database countersql.Configuration // Configuration information required to connect to the database, using the countersql package + uniqueVisits int // In memory cache of unique visitors ) // Application Startup @@ -59,7 +67,6 @@ func init() { // HTTP Routing func main() { - // API Handlers http.HandleFunc("/", countHandler) @@ -68,10 +75,14 @@ func main() { } +// HTTP handler function func countHandler(w http.ResponseWriter, r *http.Request) { if r.Method == "GET" { + // CORS header change required w.Header().Set("Access-Control-Allow-Origin", "*") w.Write([]byte(strconv.Itoa(uniqueVisits))) + + // Connect to database dbConn, err := database.Connect() if err != nil { log.Printf("failed to connect to database: %v", err) @@ -80,16 +91,17 @@ func countHandler(w http.ResponseWriter, r *http.Request) { } defer dbConn.DB.Close() + // We now get the source IP address of this request var ipAddress string - + // Check if we're behind a reverse proxy / WAF if len(r.Header.Get("X-Forwarded-For")) > 0 { ipAddress = r.Header.Get("X-Forwarded-For") } else { ipAddress = r.RemoteAddr } - ipAddress = strings.Split(ipAddress, ":")[0] + // Check if this is the first time this IP address has visited returnVisitor, err := dbConn.HasIPVisited(ipAddress) if err != nil { log.Printf("failed to determine if this is a return visitor, no data is being logged: %v", err) @@ -97,9 +109,11 @@ func countHandler(w http.ResponseWriter, r *http.Request) { } if returnVisitor { + // Log their visit err = dbConn.IncrementVisitor(ipAddress) log.Printf("return visitor from %s", ipAddress) } else { + // Insert a new visitor row in the database err = dbConn.AddVisitor(ipAddress) uniqueVisits++ log.Printf("new visitor from %s", ipAddress) @@ -110,6 +124,7 @@ func countHandler(w http.ResponseWriter, r *http.Request) { } } else { + // Needs to be GET method w.WriteHeader(http.StatusMethodNotAllowed) } }