Skip to content

feat: update refill workflow for identity-based credits #4194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Flo4604 wants to merge 6 commits into from
Closed

feat: update refill workflow for identity-based credits #4194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
feedf64
feat: update refill workflow for identity-based credits
Flo4604 Oct 28, 2025
958eb26
fix: refill new and legacy
Flo4604 Oct 30, 2025
d1c62c4
fix: refill on last day
Flo4604 Oct 30, 2025
849935c
[autofix.ci] apply automated fixes
autofix-ci[bot] Oct 30, 2025
4296f62
fix: make ts happy
Flo4604 Nov 3, 2025
768e60b
[autofix.ci] apply automated fixes
autofix-ci[bot] Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 11 additions & 34 deletions apps/docs/docs.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,9 +79,7 @@
},
{
"group": "Identities",
"pages": [
"quickstart/identities/shared-ratelimits"
]
"pages": ["quickstart/identities/shared-ratelimits"]
}
]
},
Expand Down Expand Up @@ -115,10 +113,7 @@
{
"group": "Identities",
"icon": "fingerprint",
"pages": [
"concepts/identities/overview",
"concepts/identities/ratelimits"
]
"pages": ["concepts/identities/overview", "concepts/identities/ratelimits"]
}
]
},
Expand All @@ -135,9 +130,7 @@
"pages": [
{
"group": "Ratelimiting",
"pages": [
"apis/features/ratelimiting/overview"
]
"pages": ["apis/features/ratelimiting/overview"]
},
"apis/features/temp-keys",
"apis/features/remaining",
Expand Down Expand Up @@ -172,10 +165,7 @@
{
"group": "Audit logs",
"icon": "scroll",
"pages": [
"audit-log/introduction",
"audit-log/types"
]
"pages": ["audit-log/introduction", "audit-log/types"]
},
{
"group": "Analytics",
Expand All @@ -197,10 +187,7 @@
{
"group": "Migrating API Keys",
"icon": "plane",
"pages": [
"migrations/introduction",
"migrations/keys"
]
"pages": ["migrations/introduction", "migrations/keys"]
}
]
},
Expand Down Expand Up @@ -282,9 +269,7 @@
},
{
"group": "Too Many Requests",
"pages": [
"errors/user/too_many_requests/query_quota_exceeded"
]
"pages": ["errors/user/too_many_requests/query_quota_exceeded"]
},
{
"group": "Unprocessable Entity",
Expand Down Expand Up @@ -390,15 +375,11 @@
},
{
"group": "Go",
"pages": [
"libraries/go/api"
]
"pages": ["libraries/go/api"]
},
{
"group": "Python",
"pages": [
"libraries/py/api"
]
"pages": ["libraries/py/api"]
}
]
},
Expand All @@ -423,15 +404,11 @@
},
{
"group": "Nuxt",
"pages": [
"libraries/nuxt/overview"
]
"pages": ["libraries/nuxt/overview"]
},
{
"group": "Rust",
"pages": [
"libraries/rs/overview"
]
"pages": ["libraries/rs/overview"]
},
{
"group": "Springboot",
Expand Down Expand Up @@ -474,4 +451,4 @@
"codeblocks": "system"
},
"theme": "maple"
}
}
251 changes: 193 additions & 58 deletions apps/workflows/src/workflows/refill_keys.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,77 +28,211 @@ export class RefillRemaining extends WorkflowEntrypoint<Env, Params> {
const BUCKET_NAME = "unkey_mutations";

// If refillDay is after last day of month, refillDay will be today.
const keys = await step.do(
"fetch keys",
async () =>
await db.query.keys.findMany({
where: (table, { isNotNull, isNull, and, gt, or, eq }) => {
const baseConditions = and(
isNull(table.deletedAtM),
isNotNull(table.refillAmount),
gt(table.refillAmount, table.remaining),
or(isNull(table.refillDay), eq(table.refillDay, today)),
);

if (today === lastDayOfMonth) {
return and(baseConditions, gt(table.refillDay, today));
}

return baseConditions;
},
with: {
workspace: true,
},
}),
);

console.info(`found ${keys.length} keys with refill set for today`);

for (const key of keys) {
await step.do(`refilling ${key.id}`, async () => {
const creditsToRefill = await step.do("fetch credits from credits table", async () => {
const { and, or, eq, isNull, isNotNull, gt, gte } = await import("@unkey/db");

// Build the refill conditions
const baseConditions = [
isNotNull(schema.credits.refillAmount),
gt(schema.credits.refillAmount, schema.credits.remaining),
];

if (today === lastDayOfMonth) {
// On last day of month: refill daily (null) OR refillDay >= today (catches impossible days like 31, 32, etc.)
const condition = or(
isNull(schema.credits.refillDay),
gte(schema.credits.refillDay, today),
);
if (condition) {
baseConditions.push(condition);
}
} else {
// Regular day: refill daily (null) OR refillDay === today
const condition = or(isNull(schema.credits.refillDay), eq(schema.credits.refillDay, today));
if (condition) {
baseConditions.push(condition);
}
}

const refillConditions = and(...baseConditions);

// Query with joins to filter out deleted keys/identities at DB level
const query = db
.select({
id: schema.credits.id,
workspaceId: schema.credits.workspaceId,
keyId: schema.credits.keyId,
identityId: schema.credits.identityId,
refillAmount: schema.credits.refillAmount,
})
.from(schema.credits)
.leftJoin(schema.keys, eq(schema.credits.keyId, schema.keys.id))
.leftJoin(schema.identities, eq(schema.credits.identityId, schema.identities.id))
.where(
and(
refillConditions,
// Filter: if it's a key credit, key must exist and not be deleted
// OR if it's an identity credit, identity must exist and not be deleted
or(
and(
isNotNull(schema.credits.keyId),
isNotNull(schema.keys.id),
isNull(schema.keys.deletedAtM),
),
and(
isNotNull(schema.credits.identityId),
isNotNull(schema.identities.id),
eq(schema.identities.deleted, false),
),
),
),
);

const results = await query;

return results.map((c) => ({
id: c.id,
workspaceId: c.workspaceId,
keyId: c.keyId,
identityId: c.identityId,
refillAmount: c.refillAmount as number,
isLegacy: false,
}));
});

// Also fetch legacy keys that still use the old remaining fields
const legacyKeysToRefill = await step.do("fetch legacy keys with refill", async () => {
const { and, or, eq, isNull, isNotNull, gt, gte } = await import("@unkey/db");

// Build the refill conditions for legacy keys
const baseConditions = [
isNotNull(schema.keys.refillAmount),
isNotNull(schema.keys.remaining),
gt(schema.keys.refillAmount, schema.keys.remaining),
];

if (today === lastDayOfMonth) {
// On last day of month: refill daily (null) OR refillDay >= today (catches impossible days like 31, 32, etc.)
const condition = or(isNull(schema.keys.refillDay), gte(schema.keys.refillDay, today));
if (condition) {
baseConditions.push(condition);
}
} else {
// Regular day: refill daily (null) OR refillDay === today
const condition = or(isNull(schema.keys.refillDay), eq(schema.keys.refillDay, today));
if (condition) {
baseConditions.push(condition);
}
}

const refillConditions = and(...baseConditions);

// Only get keys that DON'T have a credits entry (to avoid double refill)
const results = await db
.select({
id: schema.keys.id,
workspaceId: schema.keys.workspaceId,
refillAmount: schema.keys.refillAmount,
})
.from(schema.keys)
.leftJoin(schema.credits, eq(schema.keys.id, schema.credits.keyId))
.where(
and(
refillConditions,
isNull(schema.keys.deletedAtM),
isNull(schema.credits.id), // Key doesn't have a credits entry
),
);

return results.map((k) => ({
id: k.id,
workspaceId: k.workspaceId,
keyId: k.id,
identityId: null,
refillAmount: k.refillAmount as number,
isLegacy: true,
}));
});

const allToRefill = [...creditsToRefill, ...legacyKeysToRefill];

for (const credit of allToRefill) {
const resourceType = credit.keyId ? "key" : "identity";
const resourceId = credit.keyId || credit.identityId;
await step.do(`refilling ${credit.id} (${resourceType})`, async () => {
await db.transaction(async (tx) => {
await tx
.update(schema.keys)
.set({
remaining: key.refillAmount,
lastRefillAt: now,
})
.where(eq(schema.keys.id, key.id));
if (credit.isLegacy) {
// Update legacy key remaining field
await tx
.update(schema.keys)
.set({
remaining: credit.refillAmount,
lastRefillAt: now,
})
.where(eq(schema.keys.id, credit.id));
} else {
// Update credits table
await tx
.update(schema.credits)
.set({
remaining: credit.refillAmount,
refilledAt: now.getTime(),
updatedAt: now.getTime(),
})
.where(eq(schema.credits.id, credit.id));
}

const auditLogId = newId("auditLog");
await tx.insert(schema.auditLog).values({
id: auditLogId,
workspaceId: key.workspaceId,
bucket: BUCKET_NAME,
bucketId: "dummy",
time: now.getTime(),
event: "key.update",
actorId: "trigger",
actorType: "system",
display: `Refilled ${key.id} to ${key.refillAmount}`,
});
await tx.insert(schema.auditLogTarget).values([
const auditLogTargets = [
{
type: "workspace",
id: key.workspaceId,
workspaceId: key.workspaceId,
id: credit.workspaceId,
workspaceId: credit.workspaceId,
bucket: BUCKET_NAME,
bucketId: "dummy",
auditLogId,
displayName: `workspace ${key.workspaceId}`,
displayName: `workspace ${credit.workspaceId}`,
},
{
];

if (credit.keyId) {
auditLogTargets.push({
type: "key",
id: key.id,
workspaceId: key.workspaceId,
id: credit.keyId,
workspaceId: credit.workspaceId,
bucket: BUCKET_NAME,
bucketId: "dummy",
auditLogId,
displayName: `key ${key.id}`,
},
]);
displayName: `key ${credit.keyId}`,
});
}

if (credit.identityId) {
auditLogTargets.push({
type: "identity",
id: credit.identityId,
workspaceId: credit.workspaceId,
bucket: BUCKET_NAME,
bucketId: "dummy",
auditLogId,
displayName: `identity ${credit.identityId}`,
});
}

await tx.insert(schema.auditLog).values({
id: auditLogId,
workspaceId: credit.workspaceId,
bucket: BUCKET_NAME,
bucketId: "dummy",
time: now.getTime(),
event: resourceType === "key" ? "key.update" : "identity.update",
actorId: "trigger",
actorType: "system",
display: `Refilled ${resourceType} ${resourceId} to ${credit.refillAmount} credits`,
});
await tx.insert(schema.auditLogTarget).values(auditLogTargets);
});
return { keyId: key.id };
return { creditId: credit.id, resourceType, resourceId };
});
}

Expand All @@ -107,7 +241,8 @@ export class RefillRemaining extends WorkflowEntrypoint<Env, Params> {
});

return {
refillKeyIds: keys.map((k) => k.id),
refillCreditIds: creditsToRefill.map((c) => c.id),
refillLegacyKeyIds: legacyKeysToRefill.map((k) => k.id),
};
}
}
Loading