feat: store deployed databases info in Supabase DB

.vscode/settings.json

@@ -1,5 +1,20 @@
 {
   "deno.enablePaths": ["supabase/functions"],
   "deno.lint": true,
-  "deno.unstable": true
+  "deno.unstable": true,
+  "[javascript]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  },
+  "[json]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  },
+  "[jsonc]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  },
+  "[typescript]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  },
+  "[typescriptreact]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  }
 }

apps/db-service/package.json

@@ -9,8 +9,9 @@
     "psql": "docker compose run --rm -i psql psql"
   },
   "dependencies": {
-    "@electric-sql/pglite": "0.2.0-alpha.9",
-    "pg-gateway": "^0.2.5-alpha.2",
+    "@electric-sql/pglite": "0.2.0",
+    "@supabase/supabase-js": "^2.45.1",
+    "pg-gateway": "0.3.0-alpha.4",
     "tar": "^7.4.3"
   },
   "devDependencies": {

apps/db-service/src/index.ts

@@ -6,8 +6,11 @@ import { createReadStream } from 'node:fs'
 import { pipeline } from 'node:stream/promises'
 import { createGunzip } from 'node:zlib'
 import { extract } from 'tar'
-import { hashMd5Password, PostgresConnection, TlsOptions } from 'pg-gateway'
+import { PostgresConnection, ScramSha256Data, TlsOptions } from 'pg-gateway'
+import { createClient } from '@supabase/supabase-js'
 
+const supabaseUrl = process.env.SUPABASE_URL ?? 'http://127.0.0.1:54321'
+const supabaseKey = process.env.SUPABASE_SERVICE_ROLE_KEY ?? ''
 const dataMount = process.env.DATA_MOUNT ?? './data'
 const s3fsMount = process.env.S3FS_MOUNT ?? './s3'
 const wildcardDomain = process.env.WILDCARD_DOMAIN ?? 'db.example.com'
@@ -32,30 +35,82 @@ function getIdFromServerName(serverName: string) {
   return id
 }
 
+const PostgresErrorCodes = {
+  ConnectionException: '08000',
+} as const
+
+function sendFatalError(connection: PostgresConnection, code: string, message: string): never {
+  connection.sendError({
+    severity: 'FATAL',
+    code,
+    message,
+  })
+  connection.socket.end()
+  throw new Error(message)
+}
+
 async function fileExists(path: string): Promise<boolean> {
   try {
-    await access(path);
-    return true;
+    await access(path)
+    return true
   } catch {
-    return false;
+    return false
   }
 }
 
+const supabase = createClient(supabaseUrl, supabaseKey)
+
 const server = net.createServer((socket) => {
   let db: PGliteInterface
 
   const connection = new PostgresConnection(socket, {
     serverVersion: '16.3 (PGlite 0.2.0)',
-    authMode: 'md5Password',
-    tls,
-    async validateCredentials(credentials) {
-      if (credentials.authMode === 'md5Password') {
-        const { hash, salt } = credentials
-        const expectedHash = await hashMd5Password('postgres', 'postgres', salt)
-        return hash === expectedHash
-      }
-      return false
+    auth: {
+      method: 'scram-sha-256',
+      async getScramSha256Data(_, { tlsInfo }) {
+        if (!tlsInfo?.sniServerName) {
+          sendFatalError(
+            connection,
+            PostgresErrorCodes.ConnectionException,
+            'sniServerName required in TLS info'
+          )
+        }
+
+        const databaseId = getIdFromServerName(tlsInfo.sniServerName)
+        const { data, error } = await supabase
+          .from('deployed_databases')
+          .select('auth_method, auth_data')
+          .eq('database_id', databaseId)
+          .single()
+
+        if (error) {
+          sendFatalError(
+            connection,
+            PostgresErrorCodes.ConnectionException,
+            `Error getting auth data for database ${databaseId}: ${error}`
+          )
+        }
+
+        if (data === null) {
+          sendFatalError(
+            connection,
+            PostgresErrorCodes.ConnectionException,
+            `Database ${databaseId} not found`
+          )
+        }
+
+        if (data.auth_method !== 'scram-sha-256') {
+          sendFatalError(
+            connection,
+            PostgresErrorCodes.ConnectionException,
+            `Unsupported auth method for database ${databaseId}: ${data.auth_method}`
+          )
+        }
+
+        return data.auth_data as ScramSha256Data
+      },
     },
+    tls,
     async onTlsUpgrade({ tlsInfo }) {
       if (!tlsInfo) {
         connection.sendError({
@@ -91,12 +146,12 @@ const server = net.createServer((socket) => {
 
       console.log(`Serving database '${databaseId}'`)
 
-      const dbPath = `${dbDir}/${databaseId}`;
+      const dbPath = `${dbDir}/${databaseId}`
 
       if (!(await fileExists(dbPath))) {
         console.log(`Database '${databaseId}' is not cached, downloading...`)
 
-        const dumpPath = `${dumpDir}/${databaseId}.tar.gz`;
+        const dumpPath = `${dumpDir}/${databaseId}.tar.gz`
 
         if (!(await fileExists(dumpPath))) {
           connection.sendError({
@@ -109,33 +164,47 @@ const server = net.createServer((socket) => {
         }
 
         // Create a directory for the database
-        await mkdir(dbPath, { recursive: true });
+        await mkdir(dbPath, { recursive: true })
 
         try {
           // Extract the .tar.gz file
-          await pipeline(
-            createReadStream(dumpPath),
-            createGunzip(),
-            extract({ cwd: dbPath })
-          );
+          await pipeline(createReadStream(dumpPath), createGunzip(), extract({ cwd: dbPath }))
         } catch (error) {
-          console.error(error);
-          await rm(dbPath, { recursive: true, force: true }); // Clean up the partially created directory
+          console.error(error)
+          await rm(dbPath, { recursive: true, force: true }) // Clean up the partially created directory
           connection.sendError({
             severity: 'FATAL',
             code: 'XX000',
             message: `Error extracting database: ${(error as Error).message}`,
-          });
-          connection.socket.end();
-          return;
+          })
+          connection.socket.end()
+          return
         }
       }
 
-      db = new PGlite(dbPath, {
+      db = new PGlite({
+        dataDir: dbPath,
+        extensions: {
+          vector,
+        },
+      })
+      await db.waitReady
+      const { rows } = await db.query("SELECT 1 FROM pg_roles WHERE rolname = 'readonly_postgres';")
+      if (rows.length === 0) {
+        await db.exec(`
+          CREATE USER readonly_postgres;
+          GRANT pg_read_all_data TO readonly_postgres;
+        `)
+      }
+      await db.close()
+      db = new PGlite({
+        dataDir: dbPath,
+        username: 'readonly_postgres',
         extensions: {
           vector,
         },
       })
+      await db.waitReady
     },
     async onStartup() {
       if (!db) {
@@ -170,12 +239,12 @@ const server = net.createServer((socket) => {
     },
   })
 
-  socket.on('end', async () => {
+  socket.on('close', async () => {
     console.log('Client disconnected')
     await db?.close()
   })
 })
 
 server.listen(5432, async () => {
   console.log('Server listening on port 5432')
-})
+})

apps/postgres-new/.env.example

@@ -1,10 +1,11 @@
 NEXT_PUBLIC_SUPABASE_ANON_KEY="<supabase-anon-key>"
 NEXT_PUBLIC_SUPABASE_URL="<supabase-api-url>"
 NEXT_PUBLIC_IS_PREVIEW=true
+NEXT_PUBLIC_WILDCARD_DOMAIN=db.example.com
 
 OPENAI_API_KEY="<openai-api-key>"
 S3_ENDPOINT=http://localhost:9000
 S3_BUCKET=test
 AWS_ACCESS_KEY_ID=minioadmin
 AWS_SECRET_ACCESS_KEY=minioadmin
-WILDCARD_DOMAIN=db.example.com
+

apps/postgres-new/app/api/databases/[id]/reset-password/route.ts

@@ -0,0 +1,71 @@
+import { NextRequest, NextResponse } from 'next/server'
+import { createClient } from '~/utils/supabase/server'
+import { createScramSha256Data } from 'pg-gateway'
+import { generateDatabasePassword } from '~/utils/generate-database-password'
+
+export type DatabaseResetPasswordResponse =
+  | {
+      success: true
+      data: {
+        password: string
+      }
+    }
+  | {
+      success: false
+      error: string
+    }
+
+export async function POST(
+  req: NextRequest,
+  { params }: { params: { id: string } }
+): Promise<NextResponse<DatabaseResetPasswordResponse>> {
+  const supabase = createClient()
+
+  const {
+    data: { user },
+  } = await supabase.auth.getUser()
+
+  if (!user) {
+    return NextResponse.json(
+      {
+        success: false,
+        error: 'Unauthorized',
+      },
+      { status: 401 }
+    )
+  }
+
+  const databaseId = params.id
+
+  const { data: existingDeployedDatabase } = await supabase
+    .from('deployed_databases')
+    .select('id')
+    .eq('database_id', databaseId)
+    .maybeSingle()
+
+  if (!existingDeployedDatabase) {
+    return NextResponse.json(
+      {
+        success: false,
+        error: `Database ${databaseId} was not found`,
+      },
+      { status: 404 }
+    )
+  }
+
+  const password = generateDatabasePassword()
+
+  await supabase
+    .from('deployed_databases')
+    .update({
+      auth_data: createScramSha256Data(password),
+    })
+    .eq('database_id', databaseId)
+
+  return NextResponse.json({
+    success: true,
+    data: {
+      password,
+    },
+  })
+}

apps/postgres-new/app/api/databases/[id]/route.ts

@@ -0,0 +1,71 @@
+import { DeleteObjectCommand, S3Client } from '@aws-sdk/client-s3'
+import { NextRequest, NextResponse } from 'next/server'
+import { createClient } from '~/utils/supabase/server'
+
+const s3Client = new S3Client({ endpoint: process.env.S3_ENDPOINT, forcePathStyle: true })
+
+export type DatabaseDeleteResponse =
+  | {
+      success: true
+    }
+  | {
+      success: false
+      error: string
+    }
+
+export async function DELETE(
+  _req: NextRequest,
+  { params }: { params: { id: string } }
+): Promise<NextResponse<DatabaseDeleteResponse>> {
+  const supabase = createClient()
+
+  const {
+    data: { user },
+  } = await supabase.auth.getUser()
+
+  if (!user) {
+    return NextResponse.json(
+      {
+        success: false,
+        error: 'Unauthorized',
+      },
+      { status: 401 }
+    )
+  }
+
+  const databaseId = params.id
+
+  const { data: existingDeployedDatabase } = await supabase
+    .from('deployed_databases')
+    .select('id')
+    .eq('database_id', databaseId)
+    .maybeSingle()
+
+  if (!existingDeployedDatabase) {
+    return NextResponse.json(
+      {
+        success: false,
+        error: `Database ${databaseId} was not found`,
+      },
+      { status: 404 }
+    )
+  }
+
+  await supabase.from('deployed_databases').delete().eq('database_id', databaseId)
+
+  const key = `dbs/${databaseId}.tar.gz`
+  try {
+    await s3Client.send(
+      new DeleteObjectCommand({
+        Bucket: process.env.S3_BUCKET,
+        Key: key,
+      })
+    )
+  } catch (error) {
+    console.error(`Error deleting S3 object ${key}:`, error)
+  }
+
+  return NextResponse.json({
+    success: true,
+  })
+}