Line data Source code
1 : use measured::FixedCardinalityLabel;
2 : use serde::{Deserialize, Serialize};
3 : use std::fmt;
4 :
5 : use crate::auth::IpPattern;
6 :
7 : use crate::intern::{BranchIdInt, EndpointIdInt, ProjectIdInt};
8 :
9 : /// Generic error response with human-readable description.
10 : /// Note that we can't always present it to user as is.
11 0 : #[derive(Debug, Deserialize)]
12 : pub struct ConsoleError {
13 : pub error: Box<str>,
14 : }
15 :
16 : /// Response which holds client's auth secret, e.g. [`crate::scram::ServerSecret`].
17 : /// Returned by the `/proxy_get_role_secret` API method.
18 18 : #[derive(Deserialize)]
19 : pub struct GetRoleSecret {
20 : pub role_secret: Box<str>,
21 : pub allowed_ips: Option<Vec<IpPattern>>,
22 : pub project_id: Option<ProjectIdInt>,
23 : }
24 :
25 : // Manually implement debug to omit sensitive info.
26 : impl fmt::Debug for GetRoleSecret {
27 0 : fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28 0 : f.debug_struct("GetRoleSecret").finish_non_exhaustive()
29 0 : }
30 : }
31 :
32 : /// Response which holds compute node's `host:port` pair.
33 : /// Returned by the `/proxy_wake_compute` API method.
34 6 : #[derive(Debug, Deserialize)]
35 : pub struct WakeCompute {
36 : pub address: Box<str>,
37 : pub aux: MetricsAuxInfo,
38 : }
39 :
40 : /// Async response which concludes the link auth flow.
41 : /// Also known as `kickResponse` in the console.
42 8 : #[derive(Debug, Deserialize)]
43 : pub struct KickSession<'a> {
44 : /// Session ID is assigned by the proxy.
45 : pub session_id: &'a str,
46 :
47 : /// Compute node connection params.
48 : #[serde(deserialize_with = "KickSession::parse_db_info")]
49 : pub result: DatabaseInfo,
50 : }
51 :
52 : impl KickSession<'_> {
53 2 : fn parse_db_info<'de, D>(des: D) -> Result<DatabaseInfo, D::Error>
54 2 : where
55 2 : D: serde::Deserializer<'de>,
56 2 : {
57 4 : #[derive(Deserialize)]
58 2 : enum Wrapper {
59 2 : // Currently, console only reports `Success`.
60 2 : // `Failure(String)` used to be here... RIP.
61 2 : Success(DatabaseInfo),
62 2 : }
63 2 :
64 2 : Wrapper::deserialize(des).map(|x| match x {
65 2 : Wrapper::Success(info) => info,
66 2 : })
67 2 : }
68 : }
69 :
70 : /// Compute node connection params.
71 56 : #[derive(Deserialize)]
72 : pub struct DatabaseInfo {
73 : pub host: Box<str>,
74 : pub port: u16,
75 : pub dbname: Box<str>,
76 : pub user: Box<str>,
77 : /// Console always provides a password, but it might
78 : /// be inconvenient for debug with local PG instance.
79 : pub password: Option<Box<str>>,
80 : pub aux: MetricsAuxInfo,
81 : }
82 :
83 : // Manually implement debug to omit sensitive info.
84 : impl fmt::Debug for DatabaseInfo {
85 0 : fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
86 0 : f.debug_struct("DatabaseInfo")
87 0 : .field("host", &self.host)
88 0 : .field("port", &self.port)
89 0 : .field("dbname", &self.dbname)
90 0 : .field("user", &self.user)
91 0 : .finish_non_exhaustive()
92 0 : }
93 : }
94 :
95 : /// Various labels for prometheus metrics.
96 : /// Also known as `ProxyMetricsAuxInfo` in the console.
97 50 : #[derive(Debug, Deserialize, Clone)]
98 : pub struct MetricsAuxInfo {
99 : pub endpoint_id: EndpointIdInt,
100 : pub project_id: ProjectIdInt,
101 : pub branch_id: BranchIdInt,
102 : #[serde(default)]
103 : pub cold_start_info: ColdStartInfo,
104 : }
105 :
106 20 : #[derive(Debug, Default, Serialize, Deserialize, Clone, Copy, FixedCardinalityLabel)]
107 : #[serde(rename_all = "snake_case")]
108 : pub enum ColdStartInfo {
109 : #[default]
110 : Unknown,
111 : /// Compute was already running
112 : Warm,
113 : #[serde(rename = "pool_hit")]
114 : #[label(rename = "pool_hit")]
115 : /// Compute was not running but there was an available VM
116 : VmPoolHit,
117 : #[serde(rename = "pool_miss")]
118 : #[label(rename = "pool_miss")]
119 : /// Compute was not running and there were no VMs available
120 : VmPoolMiss,
121 :
122 : // not provided by control plane
123 : /// Connection available from HTTP pool
124 : HttpPoolHit,
125 : /// Cached connection info
126 : WarmCached,
127 : }
128 :
129 : impl ColdStartInfo {
130 0 : pub fn as_str(&self) -> &'static str {
131 0 : match self {
132 0 : ColdStartInfo::Unknown => "unknown",
133 0 : ColdStartInfo::Warm => "warm",
134 0 : ColdStartInfo::VmPoolHit => "pool_hit",
135 0 : ColdStartInfo::VmPoolMiss => "pool_miss",
136 0 : ColdStartInfo::HttpPoolHit => "http_pool_hit",
137 0 : ColdStartInfo::WarmCached => "warm_cached",
138 : }
139 0 : }
140 : }
141 :
142 : #[cfg(test)]
143 : mod tests {
144 : use super::*;
145 : use serde_json::json;
146 :
147 10 : fn dummy_aux() -> serde_json::Value {
148 10 : json!({
149 10 : "endpoint_id": "endpoint",
150 10 : "project_id": "project",
151 10 : "branch_id": "branch",
152 10 : "cold_start_info": "unknown",
153 10 : })
154 10 : }
155 :
156 : #[test]
157 2 : fn parse_kick_session() -> anyhow::Result<()> {
158 2 : // This is what the console's kickResponse looks like.
159 2 : let json = json!({
160 2 : "session_id": "deadbeef",
161 2 : "result": {
162 2 : "Success": {
163 2 : "host": "localhost",
164 2 : "port": 5432,
165 2 : "dbname": "postgres",
166 2 : "user": "john_doe",
167 2 : "password": "password",
168 2 : "aux": dummy_aux(),
169 2 : }
170 2 : }
171 2 : });
172 2 : let _: KickSession = serde_json::from_str(&json.to_string())?;
173 :
174 2 : Ok(())
175 2 : }
176 :
177 : #[test]
178 2 : fn parse_db_info() -> anyhow::Result<()> {
179 : // with password
180 2 : let _: DatabaseInfo = serde_json::from_value(json!({
181 2 : "host": "localhost",
182 2 : "port": 5432,
183 2 : "dbname": "postgres",
184 2 : "user": "john_doe",
185 2 : "password": "password",
186 2 : "aux": dummy_aux(),
187 2 : }))?;
188 :
189 : // without password
190 2 : let _: DatabaseInfo = serde_json::from_value(json!({
191 2 : "host": "localhost",
192 2 : "port": 5432,
193 2 : "dbname": "postgres",
194 2 : "user": "john_doe",
195 2 : "aux": dummy_aux(),
196 2 : }))?;
197 :
198 : // new field (forward compatibility)
199 2 : let _: DatabaseInfo = serde_json::from_value(json!({
200 2 : "host": "localhost",
201 2 : "port": 5432,
202 2 : "dbname": "postgres",
203 2 : "user": "john_doe",
204 2 : "project": "hello_world",
205 2 : "N.E.W": "forward compatibility check",
206 2 : "aux": dummy_aux(),
207 2 : }))?;
208 :
209 2 : Ok(())
210 2 : }
211 :
212 : #[test]
213 2 : fn parse_wake_compute() -> anyhow::Result<()> {
214 2 : let json = json!({
215 2 : "address": "0.0.0.0",
216 2 : "aux": dummy_aux(),
217 2 : });
218 2 : let _: WakeCompute = serde_json::from_str(&json.to_string())?;
219 2 : Ok(())
220 2 : }
221 :
222 : #[test]
223 2 : fn parse_get_role_secret() -> anyhow::Result<()> {
224 2 : // Empty `allowed_ips` field.
225 2 : let json = json!({
226 2 : "role_secret": "secret",
227 2 : });
228 2 : let _: GetRoleSecret = serde_json::from_str(&json.to_string())?;
229 2 : let json = json!({
230 2 : "role_secret": "secret",
231 2 : "allowed_ips": ["8.8.8.8"],
232 2 : });
233 2 : let _: GetRoleSecret = serde_json::from_str(&json.to_string())?;
234 2 : let json = json!({
235 2 : "role_secret": "secret",
236 2 : "allowed_ips": ["8.8.8.8"],
237 2 : "project_id": "project",
238 2 : });
239 2 : let _: GetRoleSecret = serde_json::from_str(&json.to_string())?;
240 :
241 2 : Ok(())
242 2 : }
243 : }
|
