diff --git a/controller/store/bandwidthLimitJournal_test.go b/controller/store/bandwidthLimitJournal_test.go new file mode 100644 index 00000000..254cff92 --- /dev/null +++ b/controller/store/bandwidthLimitJournal_test.go @@ -0,0 +1,63 @@ +package store + +import ( + "github.com/openziti/zrok/sdk/golang/sdk" + "github.com/stretchr/testify/assert" + "testing" +) + +func TestBandwidthLimitJournal(t *testing.T) { + str, err := Open(&Config{Path: ":memory:", Type: "sqlite3"}) + assert.NoError(t, err) + assert.NotNil(t, str) + + trx, err := str.Begin() + assert.NoError(t, err) + assert.NotNil(t, trx) + + jEmpty, err := str.IsBandwidthLimitJournalEmpty(1, trx) + assert.NoError(t, err) + assert.True(t, jEmpty) + + acctId, err := str.CreateAccount(&Account{Email: "nobody@nowhere.com", Salt: "salt", Password: "password", Token: "token"}, trx) + assert.NoError(t, err) + + _, err = str.CreateBandwidthLimitJournalEntry(&BandwidthLimitJournalEntry{AccountId: acctId, Action: WarningLimitAction, RxBytes: 1024, TxBytes: 2048}, trx) + assert.NoError(t, err) + + jEmpty, err = str.IsBandwidthLimitJournalEmpty(acctId, trx) + assert.NoError(t, err) + assert.False(t, jEmpty) + + latestJe, err := str.FindLatestBandwidthLimitJournal(acctId, trx) + assert.NoError(t, err) + assert.NotNil(t, latestJe) + assert.Nil(t, latestJe.LimitClassId) + assert.Equal(t, WarningLimitAction, latestJe.Action) + assert.Equal(t, int64(1024), latestJe.RxBytes) + assert.Equal(t, int64(2048), latestJe.TxBytes) + + lcId, err := str.CreateLimitClass(&LimitClass{ + LimitScope: AccountLimitScope, + LimitAction: LimitLimitAction, + ShareMode: sdk.PrivateShareMode, + BackendMode: sdk.VpnBackendMode, + PeriodMinutes: 60, + RxBytes: 4096, + TxBytes: 8192, + TotalBytes: 10240, + }, trx) + assert.NoError(t, err) + + _, err = str.CreateBandwidthLimitJournalEntry(&BandwidthLimitJournalEntry{AccountId: acctId, LimitClassId: &lcId, Action: LimitLimitAction, RxBytes: 10240, TxBytes: 20480}, trx) + assert.NoError(t, err) + + latestJe, err = str.FindLatestBandwidthLimitJournal(acctId, trx) + assert.NoError(t, err) + assert.NotNil(t, latestJe) + assert.NotNil(t, latestJe.LimitClassId) + assert.Equal(t, lcId, *latestJe.LimitClassId) + assert.Equal(t, LimitLimitAction, latestJe.Action) + assert.Equal(t, int64(10240), latestJe.RxBytes) + assert.Equal(t, int64(20480), latestJe.TxBytes) +} diff --git a/controller/store/limitClass.go b/controller/store/limitClass.go index 736bb44f..c5459868 100644 --- a/controller/store/limitClass.go +++ b/controller/store/limitClass.go @@ -37,7 +37,7 @@ func (str *Store) CreateLimitClass(lc *LimitClass, trx *sqlx.Tx) (int, error) { return 0, errors.Wrap(err, "error preparing limit_classes insert statement") } var id int - if err := stmt.QueryRow(lc.LimitScope, lc.LimitAction, lc.ShareMode, lc.BackendMode, lc.PeriodMinutes).Scan(&id); err != nil { + if err := stmt.QueryRow(lc.LimitScope, lc.LimitAction, lc.ShareMode, lc.BackendMode, lc.PeriodMinutes, lc.RxBytes, lc.TxBytes, lc.TotalBytes).Scan(&id); err != nil { return 0, errors.Wrap(err, "error executing limit_classes insert statement") } return id, nil diff --git a/controller/store/sql/sqlite3/022_v0_4_31_limits_classes.sql b/controller/store/sql/sqlite3/022_v0_4_31_limits_classes.sql index c8de0c00..b88cf8a5 100644 --- a/controller/store/sql/sqlite3/022_v0_4_31_limits_classes.sql +++ b/controller/store/sql/sqlite3/022_v0_4_31_limits_classes.sql @@ -1,7 +1,7 @@ -- +migrate Up create table limit_classes ( - id serial primary key, + id integer primary key, limit_scope string not null default ('account'), limit_action string not null default ('limit'), @@ -23,7 +23,7 @@ create table limit_classes ( ); create table applied_limit_classes ( - id serial primary key, + id integer primary key, account_id integer not null references accounts (id), limit_class_id integer not null references limit_classes (id), created_at datetime not null default(strftime('%Y-%m-%d %H:%M:%f', 'now')), diff --git a/controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql b/controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql index ffa03abb..135a3478 100644 --- a/controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql +++ b/controller/store/sql/sqlite3/025_v0_4_31_bandwidth_limit_journal.sql @@ -5,14 +5,14 @@ drop table environment_limit_journal; drop table share_limit_journal; create table bandwidth_limit_journal ( - id serial primary key, + id integer primary key, account_id integer references accounts (id) not null, - limit_class integer references limit_classes, + limit_class_id integer references limit_classes, action string not null, rx_bytes bigint not null, tx_bytes bigint not null, - created_at timestamptz not null default(strftime('%Y-%m-%d %H:%M:%f', 'now')), - updated_at timestamptz not null default(strftime('%Y-%m-%d %H:%M:%f', 'now')) + created_at datetime not null default(strftime('%Y-%m-%d %H:%M:%f', 'now')), + updated_at datetime not null default(strftime('%Y-%m-%d %H:%M:%f', 'now')) ); create index bandwidth_limit_journal_account_id_idx on bandwidth_limit_journal (account_id); \ No newline at end of file