Skip to content

Commit a0091d0

Browse files
guybrushguybrush
committed
(BIDS-2872) wip
1 parent 4d85feb commit a0091d0

File tree

1 file changed

+97
-55
lines changed

1 file changed

+97
-55
lines changed

‎ratelimit/ratelimit.go

+97-55
Original file line numberDiff line numberDiff line change
@@ -948,19 +948,19 @@ func DBUpdater() {
948948
func DBGetUserApiRateLimit(userId int64) (*RateLimit, error) {
949949
rl := &RateLimit{}
950950
err := db.FrontendWriterDB.Get(rl, `
951-
select second, hour, month
952-
from api_ratelimits
953-
where user_id = $1`, userId)
951+
select second, hour, month
952+
from api_ratelimits
953+
where user_id = $1`, userId)
954954
return rl, err
955955
}
956956

957957
func DBGetCurrentApiProducts() ([]*ApiProduct, error) {
958958
apiProducts := []*ApiProduct{}
959959
err := db.FrontendWriterDB.Select(&apiProducts, `
960-
select distinct on (name) name, stripe_price_id, second, hour, month, valid_from
961-
from api_products
962-
where valid_from <= now()
963-
order by name, valid_from desc`)
960+
select distinct on (name) name, stripe_price_id, second, hour, month, valid_from
961+
from api_products
962+
where valid_from <= now()
963+
order by name, valid_from desc`)
964964
return apiProducts, err
965965
}
966966

@@ -988,6 +988,17 @@ func DBUpdate(redisClient *redis.Client) error {
988988
}
989989
logrus.WithField("duration", time.Since(start)).WithField("updates", ra).Infof("updated api_ratelimits")
990990

991+
start = time.Now()
992+
res, err = DBUpdateUnlimitedRatelimits()
993+
if err != nil {
994+
return fmt.Errorf("error updating unlikmited api_ratelimit: %w", err)
995+
}
996+
ra, err = res.RowsAffected()
997+
if err != nil {
998+
return err
999+
}
1000+
logrus.WithField("duration", time.Since(start)).WithField("updates", ra).Infof("updated unlimited api_ratelimits")
1001+
9911002
start = time.Now()
9921003
res, err = DBInvalidateApiKeys()
9931004
if err != nil {
@@ -1009,65 +1020,96 @@ func DBUpdate(redisClient *redis.Client) error {
10091020
return nil
10101021
}
10111022

1012-
// DBInvalidateApiKeys invalidates api-keys where there is no corresponding user
10131023
func DBInvalidateApiKeys() (sql.Result, error) {
10141024
return db.FrontendWriterDB.Exec(`
1015-
update api_keys
1016-
set changed_at = now(), valid_until = now()
1017-
where valid_until > now() and not exists (select api_key from users where id = api_keys.user_id)`)
1025+
update api_keys
1026+
set changed_at = now(), valid_until = now()
1027+
where valid_until > now() and not exists (select api_key from users where id = api_keys.user_id)`)
10181028
}
10191029

10201030
func DBUpdateApiKeys() (sql.Result, error) {
10211031
return db.FrontendWriterDB.Exec(
10221032
`insert into api_keys (user_id, api_key, valid_until, changed_at)
1023-
select
1024-
id as user_id,
1025-
api_key,
1026-
to_timestamp('3000-01-01', 'YYYY-MM-DD') as valid_until,
1027-
now() as changed_at
1028-
from users
1029-
where api_key is not null and not exists (select user_id from api_keys where api_keys.user_id = users.id)
1030-
on conflict (user_id, api_key) do update set
1031-
valid_until = excluded.valid_until,
1032-
changed_at = excluded.changed_at
1033-
where api_keys.valid_until != excluded.valid_until`,
1033+
select
1034+
id as user_id,
1035+
api_key,
1036+
to_timestamp('3000-01-01', 'YYYY-MM-DD') as valid_until,
1037+
now() as changed_at
1038+
from users
1039+
where api_key is not null and not exists (select user_id from api_keys where api_keys.user_id = users.id)
1040+
on conflict (user_id, api_key) do update set
1041+
valid_until = excluded.valid_until,
1042+
changed_at = excluded.changed_at
1043+
where api_keys.valid_until != excluded.valid_until`,
10341044
)
10351045
}
10361046

10371047
func DBUpdateApiRatelimits() (sql.Result, error) {
10381048
return db.FrontendWriterDB.Exec(
10391049
`with
1040-
current_api_products as (
1041-
select distinct on (name) name, stripe_price_id, second, hour, month, valid_from
1042-
from api_products
1043-
where valid_from <= now()
1044-
order by name, valid_from desc
1045-
)
1046-
insert into api_ratelimits (user_id, second, hour, month, valid_until, changed_at)
1047-
select
1048-
u.id as user_id,
1049-
greatest(coalesce(cap1.second,0),coalesce(cap2.second,0)) as second,
1050-
greatest(coalesce(cap1.hour ,0),coalesce(cap2.hour ,0)) as hour,
1051-
greatest(coalesce(cap1.month ,0),coalesce(cap2.month ,0)) as month,
1052-
to_timestamp('3000-01-01', 'YYYY-MM-DD') as valid_until,
1053-
now() as changed_at
1054-
from users u
1055-
left join users_stripe_subscriptions uss on uss.customer_id = u.stripe_customer_id and uss.active = true
1056-
left join current_api_products cap on cap.stripe_price_id = uss.price_id
1057-
left join current_api_products cap1 on cap1.name = coalesce(cap.name,'free')
1058-
left join app_subs_view asv on asv.user_id = u.id and asv.active = true
1059-
left join current_api_products cap2 on cap2.name = coalesce(asv.product_id,'free')
1060-
left join api_ratelimits ar on ar.user_id = u.id
1061-
where
1062-
cap1.name != 'free' or cap2.name != 'free' or ar.user_id is not null
1063-
on conflict (user_id) do update set
1064-
second = excluded.second,
1065-
hour = excluded.hour,
1066-
month = excluded.month,
1067-
valid_until = excluded.valid_until,
1068-
changed_at = now()
1069-
where
1070-
api_ratelimits.second != excluded.second
1071-
or api_ratelimits.hour != excluded.hour
1072-
or api_ratelimits.month != excluded.month`)
1050+
current_api_products as (
1051+
select distinct on (name) name, stripe_price_id, second, hour, month, valid_from
1052+
from api_products
1053+
where valid_from <= now()
1054+
order by name, valid_from desc
1055+
)
1056+
insert into api_ratelimits (user_id, second, hour, month, valid_until, changed_at)
1057+
select
1058+
u.id as user_id,
1059+
greatest(coalesce(cap1.second,0),coalesce(cap2.second,0)) as second,
1060+
greatest(coalesce(cap1.hour ,0),coalesce(cap2.hour ,0)) as hour,
1061+
greatest(coalesce(cap1.month ,0),coalesce(cap2.month ,0)) as month,
1062+
to_timestamp('3000-01-01', 'YYYY-MM-DD') as valid_until,
1063+
now() as changed_at
1064+
from users u
1065+
left join users_stripe_subscriptions uss on uss.customer_id = u.stripe_customer_id and uss.active = true
1066+
left join current_api_products cap on cap.stripe_price_id = uss.price_id
1067+
left join current_api_products cap1 on cap1.name = coalesce(cap.name,'free')
1068+
left join app_subs_view asv on asv.user_id = u.id and asv.active = true
1069+
left join current_api_products cap2 on cap2.name = coalesce(asv.product_id,'free')
1070+
left join api_ratelimits ar on ar.user_id = u.id
1071+
where
1072+
cap1.name != 'free' or cap2.name != 'free' or ar.user_id is not null
1073+
on conflict (user_id) do update set
1074+
second = excluded.second,
1075+
hour = excluded.hour,
1076+
month = excluded.month,
1077+
valid_until = excluded.valid_until,
1078+
changed_at = now()
1079+
where
1080+
api_ratelimits.second != excluded.second
1081+
or api_ratelimits.hour != excluded.hour
1082+
or api_ratelimits.month != excluded.month`)
1083+
}
1084+
1085+
func DBUpdateUnlimitedRatelimits() (sql.Result, error) {
1086+
return db.FrontendWriterDB.Exec(
1087+
`with
1088+
unlimited_ratelimit as (
1089+
select second, hour, month
1090+
from api_products
1091+
where name = 'unlimited' and valid_from <= now()
1092+
order by valid_from desc
1093+
limit 1
1094+
)
1095+
insert into api_ratelimits (user_id, second, hour, month, valid_until, changed_at)
1096+
select
1097+
id as user_id,
1098+
unlimited_ratelimit.second,
1099+
unlimited_ratelimit.hour,
1100+
unlimited_ratelimit.month,
1101+
to_timestamp('3000-01-01', 'YYYY-MM-DD') as valid_until,
1102+
now() as changed_at
1103+
from unlimited_ratelimit, users
1104+
where user_group = 'ADMIN'
1105+
on conflict (user_id) do update set
1106+
second = excluded.second,
1107+
hour = excluded.hour,
1108+
month = excluded.month,
1109+
valid_until = excluded.valid_until,
1110+
changed_at = now()
1111+
where
1112+
api_ratelimits.second != excluded.second
1113+
or api_ratelimits.hour != excluded.hour
1114+
or api_ratelimits.month != excluded.month`)
10731115
}

0 commit comments

Comments
 (0)