From da6419a45a40715757066ee36dbefb3e2b8d35b0 Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conrad@neon.tech>
Date: Thu, 17 Jul 2025 16:20:06 +0100
Subject: [PATCH] expose lakebase-v1 as a flag

---
 proxy/src/auth/backend/mod.rs         |  5 +++++
 proxy/src/binary/proxy.rs             | 22 +++++++++++++++++++---
 proxy/src/control_plane/client/mod.rs |  6 ++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs
index 30229cec23..e865f8079c 100644
--- a/proxy/src/auth/backend/mod.rs
+++ b/proxy/src/auth/backend/mod.rs
@@ -74,6 +74,11 @@ impl std::fmt::Display for Backend<'_, ()> {
                     .debug_tuple("ControlPlane::ProxyV1")
                     .field(&endpoint.url())
                     .finish(),
+                ControlPlaneClient::LakebaseV1(lb) => fmt
+                    .debug_tuple("ControlPlane::LakebaseV1")
+                    .field(&lb.namespace)
+                    .field(&lb.port)
+                    .finish(),
                 #[cfg(any(test, feature = "testing"))]
                 ControlPlaneClient::PostgresMock(endpoint) => {
                     let url = endpoint.url();
diff --git a/proxy/src/binary/proxy.rs b/proxy/src/binary/proxy.rs
index 0ea5a89945..5abf237738 100644
--- a/proxy/src/binary/proxy.rs
+++ b/proxy/src/binary/proxy.rs
@@ -5,9 +5,7 @@ use std::pin::pin;
 use std::sync::Arc;
 use std::time::Duration;
 
-#[cfg(any(test, feature = "testing"))]
-use anyhow::Context;
-use anyhow::{bail, ensure};
+use anyhow::{Context, bail, ensure};
 use arc_swap::ArcSwapOption;
 #[cfg(any(test, feature = "testing"))]
 use camino::Utf8PathBuf;
@@ -39,6 +37,7 @@ use crate::config::{
     ProxyConfig, ProxyProtocolV2, remote_storage_from_toml,
 };
 use crate::context::parquet::ParquetUploadArgs;
+use crate::control_plane::client::lakebase_v1::LakebaseClient;
 use crate::http::health_server::AppMetrics;
 use crate::metrics::{Metrics, ServiceInfo};
 use crate::rate_limiter::{EndpointRateLimiter, RateBucketInfo, WakeComputeRateLimiter};
@@ -66,6 +65,9 @@ enum AuthBackendType {
     #[clap(alias("cplane-v1"))]
     ControlPlane,
 
+    #[clap(alias("lakebase-v1"))]
+    Lakebase,
+
     #[clap(alias("link"))]
     ConsoleRedirect,
 
@@ -734,6 +736,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
 
     match &args.auth_backend {
         AuthBackendType::ControlPlane => {}
+        AuthBackendType::Lakebase => {}
         #[cfg(any(test, feature = "testing"))]
         AuthBackendType::Postgres => {}
         #[cfg(any(test, feature = "testing"))]
@@ -828,6 +831,19 @@ fn build_auth_backend(
             Ok(Either::Left(config))
         }
 
+        AuthBackendType::Lakebase => {
+            let url: url::Url = args.auth_endpoint.parse()?;
+            let namespace = url.host_str().context("missing hostname as namespace")?;
+            let port = url.port().unwrap_or(5432);
+
+            let api = LakebaseClient::new(namespace.to_owned(), port);
+            let api = control_plane::client::ControlPlaneClient::LakebaseV1(api);
+            let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ());
+            let config = Box::leak(Box::new(auth_backend));
+
+            Ok(Either::Left(config))
+        }
+
         #[cfg(any(test, feature = "testing"))]
         AuthBackendType::Postgres => {
             let mut url: ApiUrl = args.auth_endpoint.parse()?;
diff --git a/proxy/src/control_plane/client/mod.rs b/proxy/src/control_plane/client/mod.rs
index be80171076..ee70773a77 100644
--- a/proxy/src/control_plane/client/mod.rs
+++ b/proxy/src/control_plane/client/mod.rs
@@ -29,6 +29,8 @@ use crate::types::EndpointId;
 pub enum ControlPlaneClient {
     /// Proxy V1 control plane API
     ProxyV1(cplane_proxy_v1::NeonControlPlaneClient),
+    /// Lakebase V1 mocked API.
+    LakebaseV1(lakebase_v1::LakebaseClient),
     /// Local mock control plane.
     #[cfg(any(test, feature = "testing"))]
     PostgresMock(mock::MockControlPlane),
@@ -47,6 +49,7 @@ impl ControlPlaneApi for ControlPlaneClient {
     ) -> Result<RoleAccessControl, errors::GetAuthInfoError> {
         match self {
             Self::ProxyV1(api) => api.get_role_access_control(ctx, endpoint, role).await,
+            Self::LakebaseV1(api) => api.get_role_access_control(ctx, endpoint, role).await,
             #[cfg(any(test, feature = "testing"))]
             Self::PostgresMock(api) => api.get_role_access_control(ctx, endpoint, role).await,
             #[cfg(test)]
@@ -64,6 +67,7 @@ impl ControlPlaneApi for ControlPlaneClient {
     ) -> Result<EndpointAccessControl, errors::GetAuthInfoError> {
         match self {
             Self::ProxyV1(api) => api.get_endpoint_access_control(ctx, endpoint, role).await,
+            Self::LakebaseV1(api) => api.get_endpoint_access_control(ctx, endpoint, role).await,
             #[cfg(any(test, feature = "testing"))]
             Self::PostgresMock(api) => api.get_endpoint_access_control(ctx, endpoint, role).await,
             #[cfg(test)]
@@ -78,6 +82,7 @@ impl ControlPlaneApi for ControlPlaneClient {
     ) -> Result<Vec<AuthRule>, errors::GetEndpointJwksError> {
         match self {
             Self::ProxyV1(api) => api.get_endpoint_jwks(ctx, endpoint).await,
+            Self::LakebaseV1(api) => api.get_endpoint_jwks(ctx, endpoint).await,
             #[cfg(any(test, feature = "testing"))]
             Self::PostgresMock(api) => api.get_endpoint_jwks(ctx, endpoint).await,
             #[cfg(test)]
@@ -92,6 +97,7 @@ impl ControlPlaneApi for ControlPlaneClient {
     ) -> Result<CachedNodeInfo, errors::WakeComputeError> {
         match self {
             Self::ProxyV1(api) => api.wake_compute(ctx, user_info).await,
+            Self::LakebaseV1(api) => api.wake_compute(ctx, user_info).await,
             #[cfg(any(test, feature = "testing"))]
             Self::PostgresMock(api) => api.wake_compute(ctx, user_info).await,
             #[cfg(test)]