diff --git a/torchci/components/DrCIButton.tsx b/torchci/components/DrCIButton.tsx
new file mode 100644
index 0000000000..ee43a084b2
--- /dev/null
+++ b/torchci/components/DrCIButton.tsx
@@ -0,0 +1,94 @@
+import { Button, CircularProgress, Tooltip } from "@mui/material";
+import { useSession } from "next-auth/react";
+import { useEffect, useState } from "react";
+
+export default function DrCIButton({
+ owner,
+ repo,
+ prNumber,
+}: {
+ owner: string;
+ repo: string;
+ prNumber: number;
+}) {
+ const session = useSession();
+ const loggedIn = session.status === "authenticated" && session.data !== null;
+ // loading, clickable, failed, rateLimited
+ const [buttonState, setButtonState] = useState("clickable");
+
+ const url = `/api/drci/drci?prNumber=${prNumber}`;
+ if (buttonState == "loading" && loggedIn) {
+ fetch(url, {
+ method: "POST",
+ body: JSON.stringify({ repo }),
+ headers: {
+ Authorization: session.data!["accessToken"],
+ "Cache-Control": "no-cache",
+ "Content-Type": "application/json",
+ },
+ }).then((res) => {
+ if (res.status == 429) {
+ setButtonState("rateLimited");
+ return;
+ }
+ if (!res.ok) {
+ setButtonState("failed");
+ return;
+ }
+ setButtonState("clickable");
+ return res.json();
+ });
+ }
+
+ useEffect(() => {
+ if (buttonState == "failed" || buttonState == "rateLimited") {
+ setTimeout(() => {
+ setButtonState("clickable");
+ }, 5000);
+ }
+ }, [buttonState]);
+
+ return (
+
+
+
+
+
+ );
+}
diff --git a/torchci/lib/clickhouse.ts b/torchci/lib/clickhouse.ts
index 5c69e1a008..f5aa4616c3 100644
--- a/torchci/lib/clickhouse.ts
+++ b/torchci/lib/clickhouse.ts
@@ -17,6 +17,14 @@ export function getClickhouseClient() {
});
}
+export function getClickhouseClientWritable() {
+ return createClient({
+ host: process.env.CLICKHOUSE_HUD_USER_URL ?? "http://localhost:8123",
+ username: process.env.CLICKHOUSE_HUD_USER_WRITE_USERNAME ?? "default",
+ password: process.env.CLICKHOUSE_HUD_USER_WRITE_PASSWORD ?? "",
+ });
+}
+
export async function queryClickhouse(
query: string,
params: Record
diff --git a/torchci/lib/rateLimit.ts b/torchci/lib/rateLimit.ts
new file mode 100644
index 0000000000..ab9a8630b0
--- /dev/null
+++ b/torchci/lib/rateLimit.ts
@@ -0,0 +1,39 @@
+// Rate limit users
+import dayjs from "dayjs";
+import { getClickhouseClientWritable, queryClickhouse } from "./clickhouse";
+
+async function checkRateLimit(user: string, key: string) {
+ const res = await queryClickhouse(
+ `
+select count() as count from misc.rate_limit
+where user = {user: String} and key = {key: String} and time_inserted > {timestamp: DateTime}`,
+ {
+ user,
+ key,
+ timestamp: dayjs()
+ .utc()
+ .subtract(1, "hour")
+ .format("YYYY-MM-DD HH:mm:ss"),
+ }
+ );
+ if (res.length == 0) {
+ return 0;
+ }
+ return res[0].count;
+}
+
+async function incrementRateLimit(user: string, key: string) {
+ await getClickhouseClientWritable().insert({
+ table: "misc.rate_limit",
+ values: [[user, key, dayjs().utc().format("YYYY-MM-DD HH:mm:ss")]],
+ });
+}
+
+export async function drCIRateLimitExceeded(user: string) {
+ const rateLimit = 10;
+ return (await checkRateLimit(user, "DrCI")) >= rateLimit;
+}
+
+export async function incrementDrCIRateLimit(user: string) {
+ return await incrementRateLimit(user, "DrCI");
+}
diff --git a/torchci/pages/[repoOwner]/[repoName]/pull/[prNumber].tsx b/torchci/pages/[repoOwner]/[repoName]/pull/[prNumber].tsx
index 0739c9f954..8c37874c88 100644
--- a/torchci/pages/[repoOwner]/[repoName]/pull/[prNumber].tsx
+++ b/torchci/pages/[repoOwner]/[repoName]/pull/[prNumber].tsx
@@ -1,4 +1,6 @@
+import { Stack } from "@mui/material";
import CommitStatus from "components/CommitStatus";
+import DrCIButton from "components/DrCIButton";
import { useSetTitle } from "components/DynamicTitle";
import ErrorBoundary from "components/ErrorBoundary";
import { useCHContext } from "components/UseClickhouseProvider";
@@ -133,16 +135,30 @@ function Page() {
}
return (