Skip to content
This repository has been archived by the owner on Sep 17, 2024. It is now read-only.

Commit

Permalink
chore: update raw queries to use explicit schemas for OpenGB 0.2 (#115)
Browse files Browse the repository at this point in the history
  • Loading branch information
NathanFlurry authored Jun 25, 2024
1 parent 01d9791 commit 28f260f
Show file tree
Hide file tree
Showing 6 changed files with 1,592 additions and 32 deletions.
11 changes: 7 additions & 4 deletions modules/friends/scripts/accept_request.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,15 @@ export async function run(
acceptedAt: Date | null;
declinedAt: Date | null;
}
const friendRequests = await tx.$queryRaw<FriendRequestRow[]>`
const friendRequests = await tx.$queryRawUnsafe<FriendRequestRow[]>(
`
SELECT "senderUserId", "targetUserId", "acceptedAt", "declinedAt"
FROM "FriendRequest"
WHERE "id" = ${req.friendRequestId}
FROM "${ctx.dbSchema}"."FriendRequest"
WHERE "id" = $1
FOR UPDATE
`;
`,
req.friendRequestId
);
const friendRequest = friendRequests[0];
if (!friendRequest) {
throw new RuntimeError("FRIEND_REQUEST_NOT_FOUND", {
Expand Down
15 changes: 9 additions & 6 deletions modules/friends/scripts/decline_request.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,15 @@ export async function run(
acceptedAt: Date | null;
declinedAt: Date | null;
}
const friendRequests = await tx.$queryRaw<FriendRequestRow[]>`
SELECT "senderUserId", "targetUserId", "acceptedAt", "declinedAt"
FROM "FriendRequest"
WHERE "id" = ${req.friendRequestId}
FOR UPDATE
`;
const friendRequests = await tx.$queryRawUnsafe<FriendRequestRow[]>(
`
SELECT "senderUserId", "targetUserId", "acceptedAt", "declinedAt"
FROM "${ctx.dbSchema}"."FriendRequest"
WHERE "id" = $1
FOR UPDATE
`,
req.friendRequestId
);
const friendRequest = friendRequests[0];
if (!friendRequest) {
throw new RuntimeError("FRIEND_REQUEST_NOT_FOUND", {
Expand Down
12 changes: 8 additions & 4 deletions modules/friends/scripts/send_request.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,16 @@ export async function run(
const row = await ctx.db.$transaction(async (tx) => {
// Validate that the users are not already friends
// TODO: Remove this `any` and replace with a proper type
const existingFriendRows = await tx.$queryRaw<any[]>`
const existingFriendRows = await tx.$queryRawUnsafe<any[]>(
`
SELECT 1
FROM "Friend"
WHERE "userIdA" = ${userIdA} OR "userIdB" = ${userIdA}
FROM "${ctx.dbSchema}"."Friend"
WHERE "userIdA" = $1 OR "userIdB" = $2
FOR UPDATE
`;
`,
userIdA,
userIdB,
);
if (existingFriendRows.length > 0) {
throw new RuntimeError("ALREADY_FRIENDS", { meta: { userIdA, userIdB } });
}
Expand Down
31 changes: 18 additions & 13 deletions modules/rate_limit/scripts/throttle.ts
Original file line number Diff line number Diff line change
Expand Up @@ -38,35 +38,40 @@ export async function run(
// `TokenBucket` is an unlogged table which are significantly faster to
// write to than regular tables, but are not durable. This is important
// because this script will be called on every request.
const rows = await ctx.db.$queryRaw<TokenBucket[]>`
WITH updated_bucket AS (
UPDATE "TokenBuckets" b
const rows = await ctx.db.$queryRawUnsafe<TokenBucket[]>(
`
WITH
"UpdateBucket" AS (
UPDATE "${ctx.dbSchema}"."TokenBuckets" b
SET
"tokens" = CASE
-- Reset the bucket and consume 1 token
WHEN now() > b."lastRefill" + make_interval(secs => ${req.period}) THEN ${
req.requests - 1
}
WHEN now() > b."lastRefill" + make_interval(secs => $4) THEN $3 - 1
-- Consume 1 token
ELSE b.tokens - 1
END,
"lastRefill" = CASE
WHEN now() > b."lastRefill" + make_interval(secs => ${req.period}) THEN now()
WHEN now() > b."lastRefill" + make_interval(secs => $4) THEN now()
ELSE b."lastRefill"
END
WHERE b."type" = ${req.type} AND b."key" = ${req.key}
WHERE b."type" = $1 AND b."key" = $2
RETURNING b."tokens", b."lastRefill"
),
inserted AS (
INSERT INTO "TokenBuckets" ("type", "key", "tokens", "lastRefill")
SELECT ${req.type}, ${req.key}, ${req.requests - 1}, now()
WHERE NOT EXISTS (SELECT 1 FROM updated_bucket)
INSERT INTO "${ctx.dbSchema}"."TokenBuckets" ("type", "key", "tokens", "lastRefill")
SELECT $1, $2, $3 - 1, now()
WHERE NOT EXISTS (SELECT 1 FROM "UpdateBucket")
RETURNING "tokens", "lastRefill"
)
SELECT * FROM updated_bucket
SELECT * FROM "UpdateBucket"
UNION ALL
SELECT * FROM inserted;
`;
`,
req.type,
req.key,
req.requests,
req.period,
);
const { tokens, lastRefill } = rows[0];

// If the bucket is empty, throw an error
Expand Down
13 changes: 8 additions & 5 deletions modules/tokens/scripts/revoke.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,18 +25,21 @@ export async function run(

// Sets revokedAt on all tokens that have not already been revoked. Returns
// wether or not each token was revoked.
const rows = await ctx.db.$queryRaw<TokenRow[]>`
const rows = await ctx.db.$queryRawUnsafe<TokenRow[]>(
`
WITH "PreUpdate" AS (
SELECT "id", "revokedAt"
FROM "Token"
WHERE "id" IN (${Prisma.join(req.tokenIds)})
FROM "${ctx.dbSchema}"."Token"
WHERE "id" = ANY($1)
)
UPDATE "Token"
UPDATE "${ctx.dbSchema}"."Token"
SET "revokedAt" = COALESCE("Token"."revokedAt", current_timestamp)
FROM "PreUpdate"
WHERE "Token"."id" = "PreUpdate"."id"
RETURNING "Token"."id" AS "id", "PreUpdate"."revokedAt" IS NOT NULL AS "alreadyRevoked"
`;
`,
req.tokenIds,
);

const updates: Record<string, TokenUpdate> = {};
for (const tokenId of req.tokenIds) {
Expand Down
Loading

0 comments on commit 28f260f

Please sign in to comment.