rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-05-30 19:40:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

[proxy] Add more labels to the pricing metrics

Browse Source
This commit is contained in:
Dmitry Ivanov
2022-12-26 22:10:28 +03:00
parent 7c7d225d98
commit c700c7db2e
12 changed files with 249 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
proxy/src/auth/backend.rs
View File

@@ -8,7 +8,9 @@ pub use console::{GetAuthInfoError, WakeComputeError};
use crate::{
auth::{self, AuthFlow, ClientCredentials},
compute, http, mgmt, stream, url,
compute,
console::messages::MetricsAuxInfo,
http, mgmt, stream, url,
waiters::{self, Waiter, Waiters},
};
use once_cell::sync::Lazy;
@@ -126,25 +128,13 @@ pub struct AuthSuccess<T> {
pub value: T,
}
impl<T> AuthSuccess<T> {
/// Very similar to [`std::option::Option::map`].
/// Maps [`AuthSuccess<T>`] to [`AuthSuccess<R>`] by applying
/// a function to a contained value.
pub fn map<R>(self, f: impl FnOnce(T) -> R) -> AuthSuccess<R> {
AuthSuccess {
reported_auth_ok: self.reported_auth_ok,
value: f(self.value),
}
}
}
/// Info for establishing a connection to a compute node.
/// This is what we get after auth succeeded, but not before!
pub struct NodeInfo {
/// Project from [`auth::ClientCredentials`].
pub project: String,
/// Compute node connection params.
pub config: compute::ConnCfg,
/// Labels for proxy's metrics.
pub aux: MetricsAuxInfo,
}
impl BackendType<'_, ClientCredentials<'_>> {
@@ -172,37 +162,34 @@ impl BackendType<'_, ClientCredentials<'_>> {
};
// TODO: find a proper way to merge those very similar blocks.
let (mut config, payload) = match self {
let (mut node, payload) = match self {
Console(endpoint, creds) if creds.project.is_none() => {
let payload = fetch_magic_payload.await?;
let mut creds = creds.as_ref();
creds.project = Some(payload.project.as_str().into());
let config = console::Api::new(endpoint, extra, &creds)
let node = console::Api::new(endpoint, extra, &creds)
.wake_compute()
.await?;
(config, payload)
(node, payload)
}
Postgres(endpoint, creds) if creds.project.is_none() => {
let payload = fetch_magic_payload.await?;
let mut creds = creds.as_ref();
creds.project = Some(payload.project.as_str().into());
let config = postgres::Api::new(endpoint, &creds).wake_compute().await?;
let node = postgres::Api::new(endpoint, &creds).wake_compute().await?;
(config, payload)
(node, payload)
}
_ => return Ok(None),
};
config.password(payload.password);
node.config.password(payload.password);
Ok(Some(AuthSuccess {
reported_auth_ok: false,
value: NodeInfo {
project: payload.project,
config,
},
value: node,
}))
}
@@ -233,10 +220,6 @@ impl BackendType<'_, ClientCredentials<'_>> {
console::Api::new(&endpoint, extra, &creds)
.handle_user(client)
.await?
.map(|config| NodeInfo {
project: creds.project.unwrap().into_owned(),
config,
})
}
Postgres(endpoint, creds) => {
info!("performing mock authentication using a local postgres instance");
@@ -245,10 +228,6 @@ impl BackendType<'_, ClientCredentials<'_>> {
postgres::Api::new(&endpoint, &creds)
.handle_user(client)
.await?
.map(|config| NodeInfo {
project: creds.project.unwrap().into_owned(),
config,
})
}
// NOTE: this auth backend doesn't use client credentials.
Link(url) => {

44
proxy/src/auth/backend/console.rs
View File

@@ -1,16 +1,16 @@
//! Cloud API V2.
use super::{AuthSuccess, ConsoleReqExtra};
use super::{AuthSuccess, ConsoleReqExtra, NodeInfo};
use crate::{
auth::{self, AuthFlow, ClientCredentials},
compute,
console::messages::{ConsoleError, GetRoleSecret, WakeCompute},
error::{io_error, UserFacingError},
http, sasl, scram,
stream::PqStream,
};
use futures::TryFutureExt;
use reqwest::StatusCode as HttpStatusCode;
use serde::Deserialize;
use std::future::Future;
use thiserror::Error;
use tokio::io::{AsyncRead, AsyncWrite};
@@ -136,24 +136,6 @@ impl UserFacingError for WakeComputeError {
}
}
/// Console's response which holds client's auth secret.
#[derive(Deserialize, Debug)]
struct GetRoleSecret {
role_secret: Box<str>,
}
/// Console's response which holds compute node's `host:port` pair.
#[derive(Deserialize, Debug)]
struct WakeCompute {
address: Box<str>,
}
/// Console's error response with human-readable description.
#[derive(Deserialize, Debug)]
struct ConsoleError {
error: Box<str>,
}
/// Auth secret which is managed by the cloud.
pub enum AuthInfo {
/// Md5 hash of user's password.
@@ -194,7 +176,7 @@ impl<'a> Api<'a> {
pub(super) async fn handle_user(
&'a self,
client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
) -> auth::Result<AuthSuccess<compute::ConnCfg>> {
) -> auth::Result<AuthSuccess<NodeInfo>> {
handle_user(client, self, Self::get_auth_info, Self::wake_compute).await
}
}
@@ -238,7 +220,7 @@ impl Api<'_> {
}
/// Wake up the compute node and return the corresponding connection info.
pub async fn wake_compute(&self) -> Result<compute::ConnCfg, WakeComputeError> {
pub async fn wake_compute(&self) -> Result<NodeInfo, WakeComputeError> {
let request_id = uuid::Uuid::new_v4().to_string();
async {
let request = self
@@ -269,7 +251,10 @@ impl Api<'_> {
.dbname(self.creds.dbname)
.user(self.creds.user);
Ok(config)
Ok(NodeInfo {
config,
aux: body.aux,
})
}
.map_err(crate::error::log_error)
.instrument(info_span!("wake_compute", id = request_id))
@@ -284,11 +269,11 @@ pub(super) async fn handle_user<'a, Endpoint, GetAuthInfo, WakeCompute>(
endpoint: &'a Endpoint,
get_auth_info: impl FnOnce(&'a Endpoint) -> GetAuthInfo,
wake_compute: impl FnOnce(&'a Endpoint) -> WakeCompute,
) -> auth::Result<AuthSuccess<compute::ConnCfg>>
) -> auth::Result<AuthSuccess<NodeInfo>>
where
Endpoint: AsRef<ClientCredentials<'a>>,
GetAuthInfo: Future<Output = Result<Option<AuthInfo>, GetAuthInfoError>>,
WakeCompute: Future<Output = Result<compute::ConnCfg, WakeComputeError>>,
WakeCompute: Future<Output = Result<NodeInfo, WakeComputeError>>,
{
let creds = endpoint.as_ref();
@@ -325,19 +310,20 @@ where
}
};
let mut config = wake_compute(endpoint).await?;
let mut node = wake_compute(endpoint).await?;
if let Some(keys) = scram_keys {
config.auth_keys(tokio_postgres::config::AuthKeys::ScramSha256(keys));
use tokio_postgres::config::AuthKeys;
node.config.auth_keys(AuthKeys::ScramSha256(keys));
}
Ok(AuthSuccess {
reported_auth_ok: false,
value: config,
value: node,
})
}
/// Parse http response body, taking status code into account.
async fn parse_body<T: for<'a> Deserialize<'a>>(
async fn parse_body<T: for<'a> serde::Deserialize<'a>>(
response: reqwest::Response,
) -> Result<T, ApiError> {
let status = response.status();

2
proxy/src/auth/backend/link.rs
View File

@@ -86,8 +86,8 @@ pub async fn handle_user(
Ok(AuthSuccess {
reported_auth_ok: true,
value: NodeInfo {
project: db_info.project,
config,
aux: db_info.aux,
},
})
}

11
proxy/src/auth/backend/postgres.rs
View File

@@ -2,7 +2,7 @@
use super::{
console::{self, AuthInfo, GetAuthInfoError, WakeComputeError},
AuthSuccess,
AuthSuccess, NodeInfo,
};
use crate::{
auth::{self, ClientCredentials},
@@ -57,7 +57,7 @@ impl<'a> Api<'a> {
pub(super) async fn handle_user(
&'a self,
client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin + Send>,
) -> auth::Result<AuthSuccess<compute::ConnCfg>> {
) -> auth::Result<AuthSuccess<NodeInfo>> {
// We reuse user handling logic from a production module.
console::handle_user(client, self, Self::get_auth_info, Self::wake_compute).await
}
@@ -103,7 +103,7 @@ impl Api<'_> {
}
/// We don't need to wake anything locally, so we just return the connection info.
pub async fn wake_compute(&self) -> Result<compute::ConnCfg, WakeComputeError> {
pub async fn wake_compute(&self) -> Result<NodeInfo, WakeComputeError> {
let mut config = compute::ConnCfg::new();
config
.host(self.endpoint.host_str().unwrap_or("localhost"))
@@ -111,7 +111,10 @@ impl Api<'_> {
.dbname(self.creds.dbname)
.user(self.creds.user);
Ok(config)
Ok(NodeInfo {
config,
aux: Default::default(),
})
}
}