Merge branch 'next' into chore/cache-eviction-strategy

jgoux · jgoux · commit a6028f838e11 · 2024-08-19T19:02:14.000+02:00
diff --git a/apps/proxy/.env.example b/apps/proxy/.env.example
@@ -11,9 +11,7 @@ CACHE_SCHEDULE_INTERVAL=1
 CACHE_TIMESTAMP_FILE=/var/lib/postgres-new/delete_cache_last_run
 # Cache time to live in hours
 CACHE_TTL=24
-DATA_MOUNT=/mnt/data
 S3FS_MOUNT=/mnt/s3
 SUPABASE_SERVICE_ROLE_KEY="<service-role-key>"
 SUPABASE_URL="<supabase-url>"
-WILDCARD_DOMAIN=db.example.com
-
+WILDCARD_DOMAIN=db.example.com
diff --git a/apps/proxy/Dockerfile b/apps/proxy/Dockerfile
@@ -2,7 +2,7 @@
 
 # Adjust NODE_VERSION as desired
 ARG NODE_VERSION=20.16.0
-FROM node:${NODE_VERSION}-bookworm as base
+FROM node:${NODE_VERSION}-bookworm AS base
 
 LABEL fly_launch_runtime="NodeJS"
 
@@ -13,7 +13,7 @@ WORKDIR /app
 ENV NODE_ENV=production
 
 # Build S3FS
-FROM base as build-s3fs
+FROM base AS build-s3fs
 
 # Install dependencies
 RUN apt-get update && \
@@ -28,7 +28,7 @@ RUN git clone https://github.com/supabase/s3fs-fuse.git --branch s3/support-endp
     make install
 
 # Build app
-FROM base as build-app
+FROM base AS build-app
 
 # Install packages needed to build node modules
 RUN apt-get update -qq && \
diff --git a/apps/proxy/README.md b/apps/proxy/README.md
@@ -84,8 +84,32 @@ To stop all Docker containers, run:
 docker compose down
 ```
 
-## Deployment
+## Deploying to fly.io
 
-The proxy is deployed on Fly.io.
+1. Create a new app if it doesn't exist
 
-A Tigris bucket is used to store the DB tarballs and the TLS certificates.
+   ```shell
+   fly apps create postgres-new-proxy
+   ```
+
+2. Set the appropriate environment variables and secrets for the app "postgres-new-proxy" (see `.env.example`) in fly.io UI.
+
+3. Allocate a dedicated IPv4 address for the app, it's required to resolve wildcard DNS queries correctly.
+
+   ```shell
+   fly ips allocate-v4
+   ```
+
+4. On Cloudflare, add a new CNAME DNS record pointing to the Fly.io app domain:
+
+   | Type | Name | Value |
+   | --- | --- | --- |
+   | CNAME | *.db | postgres-new-proxy.fly.dev |
+
+5. Deploy the app
+   
+   Due to a [bug](https://github.com/superfly/flyctl/issues/3870) in `fly` ignoring `.dockerignore` files when setting up a build context, we need to copy the file to the root of the repo and remove it after the deploy.
+
+   ```shell
+   cp .dockerignore ../.. && fly deploy ../.. --config apps/proxy/fly.toml && rm ../../.dockerignore
+   ```
diff --git a/apps/proxy/fly.toml b/apps/proxy/fly.toml
@@ -1,6 +1,10 @@
 app = 'postgres-new-proxy'
 
-primary_region = 'yyz'
+primary_region = 'iad'
+
+[build]
+dockerfile = "Dockerfile"
+ignorefile = ".dockerignore"
 
 [env]
 CACHE_DISK_USAGE_THRESHOLD = "90"
diff --git a/apps/proxy/package.json b/apps/proxy/package.json
@@ -10,7 +10,7 @@
     "type-check": "tsc --noEmit"
   },
   "dependencies": {
-    "@electric-sql/pglite": "0.2.0",
+    "@electric-sql/pglite": "0.2.2",
     "@supabase/supabase-js": "^2.45.1",
     "find-up": "^7.0.0",
     "pg-gateway": "0.3.0-alpha.6",
diff --git a/apps/proxy/tools/certbot/.env.example b/apps/proxy/tools/certbot/.env.example
@@ -1,7 +1,7 @@
-AWS_ACCESS_KEY_ID=625729a08b95bf1b7ff351a663f3a23c
+AWS_ACCESS_KEY_ID="<aws-access-key-id>"
 AWS_ENDPOINT_URL_S3=http://host.docker.internal:54321/storage/v1/s3
 AWS_S3_BUCKET=s3fs
-AWS_SECRET_ACCESS_KEY=850181e4652dd023b7a98c58ae0d2d34bd487ee0cc3254aed6eda37307425907
+AWS_SECRET_ACCESS_KEY="<aws-secret-access-key>"
 AWS_REGION=local
 CERTBOT_DOMAIN=db.postgres.new
 CERTBOT_EMAIL="<your-email>"
diff --git a/apps/proxy/tools/certbot/Dockerfile b/apps/proxy/tools/certbot/Dockerfile
@@ -19,7 +19,7 @@ RUN apk add --no-cache \
   fuse-dev \
   curl-dev
 
-RUN git clone https://github.com/s3fs-fuse/s3fs-fuse.git --branch v1.94 && \
+RUN git clone https://github.com/supabase/s3fs-fuse.git --depth 1 --branch s3/support-endpoint-with-uris && \
   cd s3fs-fuse && \
   ./autogen.sh && \
   ./configure && \
diff --git a/apps/proxy/tools/certbot/README.md b/apps/proxy/tools/certbot/README.md
@@ -2,7 +2,7 @@
 
 This service is responsible for managing the certificates for the PGLite instances.
 
-It uses `fly machine run --schedule weekly` to wake up the service every week to renew the certificates if needed. Let's Encrypt certificates are valid for 90 days.
+The certificates are created on Let's Encrypt (valid for 90 days) and stored in a S3 bucket under the `/tls` path.
 
 ## Testing certbot-service locally
 
@@ -32,74 +32,23 @@ The certificates will be generated in `/mnt/s3/tls`.
 
 1. Create a new app if it doesn't exist
 
-```shell
-flyctl apps create postgres-new-certbot
-```
+  ```shell
+  fly apps create postgres-new-certbot
+  ```
 
 2. Build and deploy the Docker image to fly.io image registry
 
-```shell
-flyctl deploy --build-only --push -a postgres-new-certbot --image-label
- latest
-```
+  ```shell
+  fly deploy --build-only --push --image-label latest
+  ```
 
-1. Set the appropriate environment variables and secrets for the app "postgres-new-certbot" (see `.env.example`) in fly.io UI.
-
-2. Setup [cron-manager](https://github.com/fly-apps/cron-manager?tab=readme-ov-file#getting-started) to run the certbot service every 2 weeks with the following `schedules.json`:
-
-```json
-[
-  {
-      "name": "postgres-new-certbot",
-      "app_name": "postgres-new-certbot",
-      "schedule": "0 0 1,15 * *",
-      "region": "ord",
-      "command": "./certbot.sh",
-      "command_timeout": 120,
-      "enabled": true,
-      "config": {
-          "metadata": {
-              "fly_process_group": "cron"
-          },
-          "auto_destroy": true,
-          "disable_machine_autostart": true,
-          "guest": {
-              "cpu_kind": "shared",
-              "cpus": 1,
-              "memory_mb": 256
-          },
-          "image": "registry.fly.io/postgres-new-certbot:latest",
-          "restart": {
-              "max_retries": 1,
-              "policy": "no"
-          }
-      }
-  }
-]
-```
+3. Set the appropriate environment variables and secrets for the app "postgres-new-certbot" (see `.env.example`) in fly.io UI.
 
-5. Test running the job by SSHing into cron-manager console
-
-Run this command in the cron-manager root folder:
-
-```shell
-flyctl ssh console
-```
+4. Deploy the machine with a schedule
 
-Once in the cron-manager instance:
+  ```shell
+  fly machine run registry.fly.io/postgres-new-certbot:latest "./certbot.sh" --region iad --schedule weekly
+  ```
 
-```shell
-cm jobs trigger 1
-```
-
-If you open the "postgres-new-certbot" live logs in fly.io UI, you should see the job being executed.
-
-6. You can check if the certificates are present in the Tigris bucket
-
-Run this command in the apps/db-instance folder:
-
-```shell
-flyctl storage dashboard
-```
+The machine will now be started weekly in order to renew the certificates.
 
-It should open the Tigris dashboard where you can check the bucket's content. The certificates should be created under `/tls`.
diff --git a/apps/proxy/tools/certbot/fly.toml b/apps/proxy/tools/certbot/fly.toml
@@ -1,3 +1,4 @@
 app = 'postgres-new-certbot'
 
-primary_region = 'yyz'
+[[vm]]
+size = "shared-cpu-1x"
diff --git a/apps/web/.env.example b/apps/web/.env.example
@@ -3,6 +3,9 @@ AWS_ENDPOINT_URL_S3=http://localhost:54321/storage/v1/s3
 AWS_S3_BUCKET=s3fs
 AWS_SECRET_ACCESS_KEY=850181e4652dd023b7a98c58ae0d2d34bd487ee0cc3254aed6eda37307425907
 AWS_REGION=local
+# Vercel KV (local Docker available)
+KV_REST_API_TOKEN="local_token"
+KV_REST_API_URL="http://localhost:8080"
 OPENAI_API_KEY="<openai-api-key>"
 NEXT_PUBLIC_IS_PREVIEW=true
 NEXT_PUBLIC_SUPABASE_ANON_KEY="<supabase-anon-key>"
diff --git a/apps/web/README.md b/apps/web/README.md
@@ -46,7 +46,16 @@ From this directory (`./apps/web`):
    ```shell
    echo 'OPENAI_API_KEY="<openai-api-key>"' >> .env.local
    ```
-5. Start Next.js development server:
+5. Start local Redis containers (used for rate limiting). Serves an API on port 8080:
+   ```shell
+   docker compose up -d
+   ```
+6. Store local KV (Redis) vars. Use these exact values:
+   ```shell
+   echo 'KV_REST_API_URL="http://localhost:8080"' >> .env.local
+   echo 'KV_REST_API_TOKEN="local_token"' >> .env.local
+   ```
+7. Start Next.js development server:
    ```shell
    npm run dev
    ```
diff --git a/apps/web/app/api/chat/route.ts b/apps/web/app/api/chat/route.ts
@@ -1,18 +1,51 @@
 import { openai } from '@ai-sdk/openai'
+import { Ratelimit } from '@upstash/ratelimit'
+import { kv } from '@vercel/kv'
 import { ToolInvocation, convertToCoreMessages, streamText } from 'ai'
 import { codeBlock } from 'common-tags'
 import { convertToCoreTools, maxMessageContext, maxRowLimit, tools } from '~/lib/tools'
+import { createClient } from '~/utils/supabase/server'
 
 // Allow streaming responses up to 30 seconds
 export const maxDuration = 30
 
+const inputTokenRateLimit = new Ratelimit({
+  redis: kv,
+  limiter: Ratelimit.fixedWindow(1000000, '30m'),
+  prefix: 'ratelimit:tokens:input',
+})
+
+const outputTokenRateLimit = new Ratelimit({
+  redis: kv,
+  limiter: Ratelimit.fixedWindow(10000, '30m'),
+  prefix: 'ratelimit:tokens:output',
+})
+
 type Message = {
   role: 'user' | 'assistant'
   content: string
   toolInvocations?: (ToolInvocation & { result: any })[]
 }
 
 export async function POST(req: Request) {
+  const supabase = createClient()
+
+  const { data, error } = await supabase.auth.getUser()
+
+  // We have middleware, so this should never happen (used for type narrowing)
+  if (error) {
+    return new Response('Unauthorized', { status: 401 })
+  }
+
+  const { user } = data
+
+  const { remaining: inputRemaining } = await inputTokenRateLimit.getRemaining(user.id)
+  const { remaining: outputRemaining } = await outputTokenRateLimit.getRemaining(user.id)
+
+  if (inputRemaining <= 0 || outputRemaining <= 0) {
+    return new Response('Rate limited', { status: 429 })
+  }
+
   const { messages }: { messages: Message[] } = await req.json()
 
   // Trim the message context sent to the LLM to mitigate token abuse
@@ -64,6 +97,14 @@ export async function POST(req: Request) {
     model: openai('gpt-4o-2024-08-06'),
     messages: convertToCoreMessages(trimmedMessageContext),
     tools: convertToCoreTools(tools),
+    async onFinish({ usage }) {
+      await inputTokenRateLimit.limit(user.id, {
+        rate: usage.promptTokens,
+      })
+      await outputTokenRateLimit.limit(user.id, {
+        rate: usage.completionTokens,
+      })
+    },
   })
 
   return result.toAIStreamResponse()
diff --git a/apps/web/app/api/databases/[id]/upload/route.ts b/apps/web/app/api/databases/[id]/upload/route.ts
@@ -6,6 +6,7 @@ import { Readable } from 'stream'
 import { createClient } from '~/utils/supabase/server'
 import { createScramSha256Data } from 'pg-gateway'
 import { generateDatabasePassword } from '~/utils/generate-database-password'
+import { getUncompressedSizeInMB } from '~/utils/get-uncompressed-size-in-mb'
 
 const wildcardDomain = process.env.NEXT_PUBLIC_WILDCARD_DOMAIN ?? 'db.example.com'
 const s3Client = new S3Client({ forcePathStyle: true })
@@ -62,8 +63,7 @@ export async function POST(
     )
   }
 
-  const dumpSizeInMB = dump.size / (1024 * 1024)
-  if (dumpSizeInMB > 100) {
+  if ((await getUncompressedSizeInMB(dump)) > 100) {
     return NextResponse.json(
       {
         success: false,
diff --git a/apps/web/components/app-provider.tsx b/apps/web/components/app-provider.tsx
@@ -27,6 +27,7 @@ const dbManager = typeof window !== 'undefined' ? new DbManager() : undefined
 export default function AppProvider({ children }: AppProps) {
   const [isLoadingUser, setIsLoadingUser] = useState(true)
   const [user, setUser] = useState<User>()
+  const [isRateLimited, setIsRateLimited] = useState(false)
 
   const focusRef = useRef<FocusHandle>(null)
 
@@ -110,6 +111,8 @@ export default function AppProvider({ children }: AppProps) {
         isLoadingUser,
         signIn,
         signOut,
+        isRateLimited,
+        setIsRateLimited,
         focusRef,
         isPreview,
         dbManager,
@@ -131,6 +134,8 @@ export type AppContextValues = {
   isLoadingUser: boolean
   signIn: () => Promise<User | undefined>
   signOut: () => Promise<void>
+  isRateLimited: boolean
+  setIsRateLimited: (limited: boolean) => void
   focusRef: RefObject<FocusHandle>
   isPreview: boolean
   dbManager?: DbManager
diff --git a/apps/web/components/chat.tsx b/apps/web/components/chat.tsx
diff --git a/apps/web/components/workspace.tsx b/apps/web/components/workspace.tsx
diff --git a/apps/web/data/deployed-databases/deployed-database-create-mutation.ts b/apps/web/data/deployed-databases/deployed-database-create-mutation.ts
diff --git a/apps/web/docker-compose.yml b/apps/web/docker-compose.yml
diff --git a/apps/web/package.json b/apps/web/package.json
diff --git a/apps/web/utils/get-uncompressed-size-in-mb.ts b/apps/web/utils/get-uncompressed-size-in-mb.ts
diff --git a/package-lock.json b/package-lock.json
