@@ -2,24 +2,81 @@ package main
22
33import (
44 "database/sql"
5- "fmt"
65 "os"
76 "time"
87
9- "github.com/coder/flog "
8+ "github.com/gocarina/gocsv "
109 "github.com/spf13/cobra"
10+ "golang.org/x/xerrors"
11+
12+ "github.com/coder/flog"
1113)
1214
1315func connectDB () (* sql.DB , error ) {
1416 const envKey = "POSTGRES_URL"
1517 url := os .Getenv (envKey )
1618 if url == "" {
17- return nil , fmt .Errorf ("no $%v provided" , envKey )
19+ return nil , xerrors .Errorf ("no $%v provided" , envKey )
1820 }
1921
2022 return sql .Open ("postgres" , url )
2123}
2224
25+ type trainingRow struct {
26+ WorkspaceID string `csv:"workspace_id"`
27+ // HourOfDay ranges from 0 to 23
28+ HourOfDay int `csv:"hour"`
29+ // Day of Week ranges from 0 to 6
30+ DayOfWeek int `csv:"day"`
31+ Used bool `csv:"used"`
32+ }
33+
34+ type dbRow struct {
35+ Time time.Time
36+ WorkspaceID string
37+ }
38+
39+ func (db dbRow ) convert (used bool ) trainingRow {
40+ return trainingRow {
41+ WorkspaceID : db .WorkspaceID ,
42+ HourOfDay : db .Time .Hour (),
43+ DayOfWeek : int (db .Time .Weekday ()),
44+ Used : used ,
45+ }
46+ }
47+
48+ // generateTrainingRows accepts sparse input data from the DB and creates
49+ // trainingRows suitable to enter a prediction model.
50+ func generateTrainingRows (rs []dbRow ) []trainingRow {
51+ workspaceIDs := make (map [string ]struct {})
52+ for _ , r := range rs {
53+ workspaceIDs [r .WorkspaceID ] = struct {}{}
54+ }
55+
56+ var trainingRows []trainingRow
57+
58+ last := rs [0 ].Time
59+ for _ , r := range rs {
60+ if ! r .Time .Equal (last ) && ! last .IsZero () {
61+ // We just skipped a time-slot, we must fill in the blanks.
62+ for last .Before (r .Time ) {
63+ last = last .Add (time .Hour )
64+ for wid := range workspaceIDs {
65+ trainingRows = append (trainingRows , trainingRow {
66+ WorkspaceID : wid ,
67+ HourOfDay : last .Hour (),
68+ DayOfWeek : int (last .Weekday ()),
69+ Used : false ,
70+ })
71+ }
72+ }
73+ }
74+ trainingRows = append (trainingRows , r .convert (true ))
75+ }
76+
77+ return trainingRows
78+ }
79+
2380func loadTrainingCSV () * cobra.Command {
2481 return & cobra.Command {
2582 Use : "load-training-csv" ,
@@ -28,28 +85,46 @@ func loadTrainingCSV() *cobra.Command {
2885 if err != nil {
2986 return err
3087 }
31- const q = `SELECT
32- date_trunc('hour', created_at) at_hour
88+ const q = `
89+ SELECT
90+ date_trunc('hour', ag.created_at) at_hour,
91+ workspace_id
3392 FROM
34- agent_stats
93+ agent_stats ag
94+ JOIN workspaces w ON
95+ w.id = ag.workspace_id
96+ WHERE
97+ NOT w.deleted
3598 GROUP BY
36- workspace_id, at_hour;
99+ workspace_id,
100+ user_id,
101+ at_hour
102+ ORDER BY
103+ at_hour ASC;
37104 `
38105 rows , err := db .Query (q )
39106 if err != nil {
40107 return err
41108 }
42109
43- var times []time. Time
110+ var rs []dbRow
44111 for rows .Next () {
45- var t time. Time
46- err = rows .Scan (& t )
112+ var r dbRow
113+ err = rows .Scan (& r . Time , & r . WorkspaceID )
47114 if err != nil {
48- return nil
115+ return err
49116 }
50- times = append (times , t )
117+ rs = append (rs , r )
51118 }
52- flog .Info ("times: %+v" , times )
119+ err = rows .Err ()
120+ if err != nil {
121+ return err
122+ }
123+
124+ flog .Info ("loaded %v rows" , len (rs ))
125+ trainingRows := generateTrainingRows (rs )
126+ flog .Info ("generated %v training rows" , len (trainingRows ))
127+ gocsv .Marshal (trainingRows , os .Stdout )
53128 return nil
54129 },
55130 }
0 commit comments