/
model.go
222 lines (191 loc) · 6.11 KB
/
model.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
package v_0005
import (
"context"
"database/sql"
"database/sql/driver"
"encoding/hex"
"encoding/json"
"time"
"github.com/google/uuid"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
type Status string
const (
StatusRunning Status = "RUNNING"
StatusScheduled Status = "SCHEDULED"
StatusFinished Status = "FINISHED"
StatusFailed Status = "FAILED"
StatusKilled Status = "KILLED"
)
type LifecycleStage string
const (
LifecycleStageActive LifecycleStage = "active"
LifecycleStageDeleted LifecycleStage = "deleted"
)
type Experiment struct {
ID *int32 `gorm:"column:experiment_id;not null;primaryKey"`
Name string `gorm:"type:varchar(256);not null;unique"`
ArtifactLocation string `gorm:"type:varchar(256)"`
LifecycleStage LifecycleStage `gorm:"type:varchar(32);check:lifecycle_stage IN ('active', 'deleted')"`
CreationTime sql.NullInt64 `gorm:"type:bigint"`
LastUpdateTime sql.NullInt64 `gorm:"type:bigint"`
Tags []ExperimentTag `gorm:"constraint:OnDelete:CASCADE"`
Runs []Run `gorm:"constraint:OnDelete:CASCADE"`
}
type ExperimentTag struct {
Key string `gorm:"type:varchar(250);not null;primaryKey"`
Value string `gorm:"type:varchar(5000)"`
ExperimentID int32 `gorm:"not null;primaryKey"`
}
//nolint:lll
type Run struct {
ID string `gorm:"<-:create;column:run_uuid;type:varchar(32);not null;primaryKey"`
Name string `gorm:"type:varchar(250)"`
SourceType string `gorm:"<-:create;type:varchar(20);check:source_type IN ('NOTEBOOK', 'JOB', 'LOCAL', 'UNKNOWN', 'PROJECT')"`
SourceName string `gorm:"<-:create;type:varchar(500)"`
EntryPointName string `gorm:"<-:create;type:varchar(50)"`
UserID string `gorm:"<-:create;type:varchar(256)"`
Status Status `gorm:"type:varchar(9);check:status IN ('SCHEDULED', 'FAILED', 'FINISHED', 'RUNNING', 'KILLED')"`
StartTime sql.NullInt64 `gorm:"<-:create;type:bigint"`
EndTime sql.NullInt64 `gorm:"type:bigint"`
SourceVersion string `gorm:"<-:create;type:varchar(50)"`
LifecycleStage LifecycleStage `gorm:"type:varchar(20);check:lifecycle_stage IN ('active', 'deleted')"`
ArtifactURI string `gorm:"<-:create;type:varchar(200)"`
ExperimentID int32
Experiment Experiment
DeletedTime sql.NullInt64 `gorm:"type:bigint"`
RowNum RowNum `gorm:"<-:create;index"`
Params []Param `gorm:"constraint:OnDelete:CASCADE"`
Tags []Tag `gorm:"constraint:OnDelete:CASCADE"`
Metrics []Metric `gorm:"constraint:OnDelete:CASCADE"`
LatestMetrics []LatestMetric `gorm:"constraint:OnDelete:CASCADE"`
}
type RowNum int64
func (rn *RowNum) Scan(v interface{}) error {
nullInt := sql.NullInt64{}
if err := nullInt.Scan(v); err != nil {
return err
}
*rn = RowNum(nullInt.Int64)
return nil
}
func (rn RowNum) GormDataType() string {
return "bigint"
}
func (rn RowNum) GormValue(ctx context.Context, db *gorm.DB) clause.Expr {
if rn == 0 {
return clause.Expr{
SQL: "(SELECT COALESCE(MAX(row_num), -1) FROM runs) + 1",
}
}
return clause.Expr{
SQL: "?",
Vars: []interface{}{int64(rn)},
}
}
type Param struct {
Key string `gorm:"type:varchar(250);not null;primaryKey"`
Value string `gorm:"type:varchar(500);not null"`
RunID string `gorm:"column:run_uuid;not null;primaryKey;index"`
}
type Tag struct {
Key string `gorm:"type:varchar(250);not null;primaryKey"`
Value string `gorm:"type:varchar(5000)"`
RunID string `gorm:"column:run_uuid;not null;primaryKey;index"`
}
type Metric struct {
Key string `gorm:"type:varchar(250);not null;primaryKey"`
Value float64 `gorm:"type:double precision;not null;primaryKey"`
Timestamp int64 `gorm:"not null;primaryKey"`
RunID string `gorm:"column:run_uuid;not null;primaryKey;index"`
Step int64 `gorm:"default:0;not null;primaryKey"`
IsNan bool `gorm:"default:false;not null;primaryKey"`
Iter int64 `gorm:"index"`
}
type LatestMetric struct {
Key string `gorm:"type:varchar(250);not null;primaryKey"`
Value float64 `gorm:"type:double precision;not null"`
Timestamp int64
Step int64 `gorm:"not null"`
IsNan bool `gorm:"not null"`
RunID string `gorm:"column:run_uuid;not null;primaryKey;index"`
LastIter int64
}
type AlembicVersion struct {
Version string `gorm:"column:version_num;type:varchar(32);not null;primaryKey"`
}
func (AlembicVersion) TableName() string {
return "alembic_version"
}
type SchemaVersion struct {
Version string `gorm:"not null;primaryKey"`
}
func (SchemaVersion) TableName() string {
return "schema_version"
}
type Base struct {
ID uuid.UUID `gorm:"type:uuid;primaryKey" json:"id"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
IsArchived bool `json:"-"`
}
func (b *Base) BeforeCreate(tx *gorm.DB) error {
b.ID = uuid.New()
return nil
}
type Dashboard struct {
Base
Name string `json:"name"`
Description string `json:"description"`
AppID *uuid.UUID `gorm:"type:uuid" json:"app_id"`
App App `json:"-"`
}
func (d Dashboard) MarshalJSON() ([]byte, error) {
type localDashboard Dashboard
type jsonDashboard struct {
localDashboard
AppType *string `json:"app_type"`
}
jd := jsonDashboard{
localDashboard: localDashboard(d),
}
if d.App.IsArchived {
jd.AppID = nil
} else {
jd.AppType = &d.App.Type
}
return json.Marshal(jd)
}
type App struct {
Base
Type string `gorm:"not null" json:"type"`
State AppState `json:"state"`
}
type AppState map[string]any
func (s AppState) Value() (driver.Value, error) {
v, err := json.Marshal(s)
if err != nil {
return nil, err
}
return string(v), nil
}
func (s *AppState) Scan(v interface{}) error {
var nullS sql.NullString
if err := nullS.Scan(v); err != nil {
return err
}
if nullS.Valid {
return json.Unmarshal([]byte(nullS.String), s)
}
return nil
}
func (s AppState) GormDataType() string {
return "text"
}
func NewUUID() string {
var r [32]byte
u := uuid.New()
hex.Encode(r[:], u[:])
return string(r[:])
}