expose lakebase-v1 as a flag

This commit is contained in:
Conrad Ludgate
2025-07-17 16:20:06 +01:00
parent e5f5c79eb1
commit da6419a45a
3 changed files with 30 additions and 3 deletions

View File

@@ -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();

View File

@@ -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()?;

View File

@@ -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)]