mirror of
https://github.com/hatchet-dev/hatchet.git
synced 2025-12-20 08:10:26 -06:00
344 lines
8.7 KiB
Go
344 lines
8.7 KiB
Go
// Code generated by sqlc. DO NOT EDIT.
|
|
// versions:
|
|
// sqlc v1.29.0
|
|
// source: rate_limits.sql
|
|
|
|
package sqlcv1
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
)
|
|
|
|
const bulkUpdateRateLimits = `-- name: BulkUpdateRateLimits :many
|
|
WITH input AS (
|
|
SELECT
|
|
"key", "units"
|
|
FROM
|
|
(
|
|
SELECT
|
|
unnest($1::text[]) AS "key",
|
|
unnest($2::int[]) AS "units"
|
|
) AS subquery
|
|
), rls_to_update AS (
|
|
SELECT
|
|
rl."tenantId", rl.key, rl."limitValue", rl.value, rl."window", rl."lastRefill"
|
|
FROM
|
|
"RateLimit" rl
|
|
WHERE
|
|
rl."tenantId" = $3::uuid
|
|
AND rl."key" = ANY(SELECT "key" FROM input)
|
|
ORDER BY
|
|
rl."tenantId" ASC, rl."key" ASC
|
|
FOR UPDATE
|
|
)
|
|
UPDATE
|
|
"RateLimit" rl
|
|
SET
|
|
"value" = get_refill_value(rl) - (SELECT "units" FROM input WHERE "key" = rl."key"),
|
|
"lastRefill" = CASE
|
|
WHEN NOW() - rl."lastRefill" >= (rl."window"::INTERVAL - INTERVAL '10 milliseconds') THEN
|
|
CURRENT_TIMESTAMP
|
|
ELSE
|
|
rl."lastRefill"
|
|
END
|
|
FROM
|
|
rls_to_update rl2
|
|
WHERE
|
|
rl2."tenantId" = rl."tenantId"
|
|
AND rl2."key" = rl."key"
|
|
RETURNING rl."tenantId", rl.key, rl."limitValue", rl.value, rl."window", rl."lastRefill"
|
|
`
|
|
|
|
type BulkUpdateRateLimitsParams struct {
|
|
Keys []string `json:"keys"`
|
|
Units []int32 `json:"units"`
|
|
Tenantid pgtype.UUID `json:"tenantid"`
|
|
}
|
|
|
|
func (q *Queries) BulkUpdateRateLimits(ctx context.Context, db DBTX, arg BulkUpdateRateLimitsParams) ([]*RateLimit, error) {
|
|
rows, err := db.Query(ctx, bulkUpdateRateLimits, arg.Keys, arg.Units, arg.Tenantid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []*RateLimit
|
|
for rows.Next() {
|
|
var i RateLimit
|
|
if err := rows.Scan(
|
|
&i.TenantId,
|
|
&i.Key,
|
|
&i.LimitValue,
|
|
&i.Value,
|
|
&i.Window,
|
|
&i.LastRefill,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, &i)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
const listRateLimitsForSteps = `-- name: ListRateLimitsForSteps :many
|
|
SELECT
|
|
units, "stepId", "rateLimitKey", "tenantId", kind
|
|
FROM
|
|
"StepRateLimit" srl
|
|
WHERE
|
|
srl."stepId" = ANY($1::uuid[])
|
|
AND srl."tenantId" = $2::uuid
|
|
`
|
|
|
|
type ListRateLimitsForStepsParams struct {
|
|
Stepids []pgtype.UUID `json:"stepids"`
|
|
Tenantid pgtype.UUID `json:"tenantid"`
|
|
}
|
|
|
|
func (q *Queries) ListRateLimitsForSteps(ctx context.Context, db DBTX, arg ListRateLimitsForStepsParams) ([]*StepRateLimit, error) {
|
|
rows, err := db.Query(ctx, listRateLimitsForSteps, arg.Stepids, arg.Tenantid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []*StepRateLimit
|
|
for rows.Next() {
|
|
var i StepRateLimit
|
|
if err := rows.Scan(
|
|
&i.Units,
|
|
&i.StepId,
|
|
&i.RateLimitKey,
|
|
&i.TenantId,
|
|
&i.Kind,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, &i)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
const listRateLimitsForTenantNoMutate = `-- name: ListRateLimitsForTenantNoMutate :many
|
|
SELECT
|
|
"tenantId",
|
|
"key",
|
|
"limitValue",
|
|
(CASE
|
|
WHEN NOW() - rl."lastRefill" >= rl."window"::INTERVAL THEN
|
|
get_refill_value(rl)
|
|
ELSE
|
|
rl."value"
|
|
END)::int AS "value",
|
|
"window",
|
|
(CASE
|
|
WHEN NOW() - rl."lastRefill" >= rl."window"::INTERVAL THEN
|
|
CURRENT_TIMESTAMP
|
|
ELSE
|
|
rl."lastRefill"
|
|
END)::timestamp AS "lastRefill"
|
|
FROM
|
|
"RateLimit" rl
|
|
WHERE
|
|
"tenantId" = $1::uuid
|
|
AND (
|
|
$2::text IS NULL OR
|
|
rl."key" like concat('%', $2::text, '%')
|
|
)
|
|
ORDER BY
|
|
rl."key" ASC
|
|
OFFSET
|
|
COALESCE($3, 0)
|
|
LIMIT
|
|
COALESCE($4, 50)
|
|
`
|
|
|
|
type ListRateLimitsForTenantNoMutateParams struct {
|
|
Tenantid pgtype.UUID `json:"tenantid"`
|
|
Search pgtype.Text `json:"search"`
|
|
Offset interface{} `json:"offset"`
|
|
Limit interface{} `json:"limit"`
|
|
}
|
|
|
|
type ListRateLimitsForTenantNoMutateRow struct {
|
|
TenantId pgtype.UUID `json:"tenantId"`
|
|
Key string `json:"key"`
|
|
LimitValue int32 `json:"limitValue"`
|
|
Value int32 `json:"value"`
|
|
Window string `json:"window"`
|
|
LastRefill pgtype.Timestamp `json:"lastRefill"`
|
|
}
|
|
|
|
// Returns the same results as ListRateLimitsForTenantWithMutate but does not update the rate limit values
|
|
func (q *Queries) ListRateLimitsForTenantNoMutate(ctx context.Context, db DBTX, arg ListRateLimitsForTenantNoMutateParams) ([]*ListRateLimitsForTenantNoMutateRow, error) {
|
|
rows, err := db.Query(ctx, listRateLimitsForTenantNoMutate,
|
|
arg.Tenantid,
|
|
arg.Search,
|
|
arg.Offset,
|
|
arg.Limit,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []*ListRateLimitsForTenantNoMutateRow
|
|
for rows.Next() {
|
|
var i ListRateLimitsForTenantNoMutateRow
|
|
if err := rows.Scan(
|
|
&i.TenantId,
|
|
&i.Key,
|
|
&i.LimitValue,
|
|
&i.Value,
|
|
&i.Window,
|
|
&i.LastRefill,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, &i)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
const listRateLimitsForTenantWithMutate = `-- name: ListRateLimitsForTenantWithMutate :many
|
|
WITH rls_to_update AS (
|
|
SELECT
|
|
rl."tenantId", rl.key, rl."limitValue", rl.value, rl."window", rl."lastRefill"
|
|
FROM
|
|
"RateLimit" rl
|
|
WHERE
|
|
rl."tenantId" = $1::uuid
|
|
AND NOW() - rl."lastRefill" >= (rl."window"::INTERVAL - INTERVAL '10 milliseconds')
|
|
ORDER BY
|
|
rl."tenantId" ASC, rl."key" ASC
|
|
FOR UPDATE
|
|
), refill AS (
|
|
UPDATE
|
|
"RateLimit" rl
|
|
SET
|
|
"value" = get_refill_value(rl),
|
|
"lastRefill" = CURRENT_TIMESTAMP
|
|
FROM
|
|
rls_to_update
|
|
WHERE
|
|
rl."tenantId" = rls_to_update."tenantId"
|
|
AND rl."key" = rls_to_update."key"
|
|
RETURNING rl."tenantId", rl.key, rl."limitValue", rl.value, rl."window", rl."lastRefill"
|
|
)
|
|
SELECT
|
|
rl."tenantId", rl.key, rl."limitValue", rl.value, rl."window", rl."lastRefill",
|
|
(rl."lastRefill" + rl."window"::INTERVAL - INTERVAL '10 milliseconds')::timestamp AS "nextRefillAt"
|
|
FROM
|
|
"RateLimit" rl
|
|
WHERE
|
|
rl."tenantId" = $1::uuid
|
|
AND rl."key" NOT IN (SELECT "key" FROM refill)
|
|
|
|
UNION ALL
|
|
|
|
SELECT
|
|
refill."tenantId", refill.key, refill."limitValue", refill.value, refill."window", refill."lastRefill",
|
|
-- return the next refill time
|
|
(refill."lastRefill" + refill."window"::INTERVAL - INTERVAL '10 milliseconds')::timestamp AS "nextRefillAt"
|
|
FROM
|
|
refill
|
|
`
|
|
|
|
type ListRateLimitsForTenantWithMutateRow struct {
|
|
TenantId pgtype.UUID `json:"tenantId"`
|
|
Key string `json:"key"`
|
|
LimitValue int32 `json:"limitValue"`
|
|
Value int32 `json:"value"`
|
|
Window string `json:"window"`
|
|
LastRefill pgtype.Timestamp `json:"lastRefill"`
|
|
NextRefillAt pgtype.Timestamp `json:"nextRefillAt"`
|
|
}
|
|
|
|
func (q *Queries) ListRateLimitsForTenantWithMutate(ctx context.Context, db DBTX, tenantid pgtype.UUID) ([]*ListRateLimitsForTenantWithMutateRow, error) {
|
|
rows, err := db.Query(ctx, listRateLimitsForTenantWithMutate, tenantid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var items []*ListRateLimitsForTenantWithMutateRow
|
|
for rows.Next() {
|
|
var i ListRateLimitsForTenantWithMutateRow
|
|
if err := rows.Scan(
|
|
&i.TenantId,
|
|
&i.Key,
|
|
&i.LimitValue,
|
|
&i.Value,
|
|
&i.Window,
|
|
&i.LastRefill,
|
|
&i.NextRefillAt,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, &i)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
const upsertRateLimitsBulk = `-- name: UpsertRateLimitsBulk :exec
|
|
WITH input_values AS (
|
|
SELECT
|
|
"key", "limitValue", "window"
|
|
FROM
|
|
(
|
|
SELECT
|
|
unnest($2::text[]) AS "key",
|
|
unnest($3::int[]) AS "limitValue",
|
|
unnest($4::text[]) AS "window"
|
|
) AS subquery
|
|
ORDER BY
|
|
"key"
|
|
)
|
|
INSERT INTO "RateLimit" (
|
|
"tenantId",
|
|
"key",
|
|
"limitValue",
|
|
"value",
|
|
"window"
|
|
)
|
|
SELECT
|
|
$1::uuid,
|
|
iv."key",
|
|
iv."limitValue",
|
|
iv."limitValue",
|
|
iv."window"
|
|
FROM
|
|
input_values iv
|
|
ON CONFLICT ("tenantId", "key") DO UPDATE SET
|
|
"limitValue" = EXCLUDED."limitValue",
|
|
"window" = EXCLUDED."window",
|
|
"value" = CASE WHEN EXCLUDED."limitValue" < "RateLimit"."value" THEN EXCLUDED."limitValue" ELSE "RateLimit"."value" END
|
|
`
|
|
|
|
type UpsertRateLimitsBulkParams struct {
|
|
Tenantid pgtype.UUID `json:"tenantid"`
|
|
Keys []string `json:"keys"`
|
|
Limitvalues []int32 `json:"limitvalues"`
|
|
Windows []string `json:"windows"`
|
|
}
|
|
|
|
func (q *Queries) UpsertRateLimitsBulk(ctx context.Context, db DBTX, arg UpsertRateLimitsBulkParams) error {
|
|
_, err := db.Exec(ctx, upsertRateLimitsBulk,
|
|
arg.Tenantid,
|
|
arg.Keys,
|
|
arg.Limitvalues,
|
|
arg.Windows,
|
|
)
|
|
return err
|
|
}
|