proxy: Record and export user-agent header (#10955)

neondatabase/cloud#24464

proxy/src/context/mod.rs

@@ -55,6 +55,7 @@ struct RequestContextInner {
     dbname: Option<DbName>,
     user: Option<RoleName>,
     application: Option<SmolStr>,
+    user_agent: Option<SmolStr>,
     error_kind: Option<ErrorKind>,
     pub(crate) auth_method: Option<AuthMethod>,
     jwt_issuer: Option<String>,
@@ -100,6 +101,7 @@ impl Clone for RequestContext {
             dbname: inner.dbname.clone(),
             user: inner.user.clone(),
             application: inner.application.clone(),
+            user_agent: inner.user_agent.clone(),
             error_kind: inner.error_kind,
             auth_method: inner.auth_method.clone(),
             jwt_issuer: inner.jwt_issuer.clone(),
@@ -149,6 +151,7 @@ impl RequestContext {
             dbname: None,
             user: None,
             application: None,
+            user_agent: None,
             error_kind: None,
             auth_method: None,
             jwt_issuer: None,
@@ -245,6 +248,13 @@ impl RequestContext {
             .set_user(user);
     }
 
+    pub(crate) fn set_user_agent(&self, user_agent: Option<SmolStr>) {
+        self.0
+            .try_lock()
+            .expect("should not deadlock")
+            .set_user_agent(user_agent);
+    }
+
     pub(crate) fn set_auth_method(&self, auth_method: AuthMethod) {
         let mut this = self.0.try_lock().expect("should not deadlock");
         this.auth_method = Some(auth_method);
@@ -384,6 +394,10 @@ impl RequestContextInner {
         }
     }
 
+    fn set_user_agent(&mut self, user_agent: Option<SmolStr>) {
+        self.user_agent = user_agent;
+    }
+
     fn set_dbname(&mut self, dbname: DbName) {
         self.dbname = Some(dbname);
     }

proxy/src/context/parquet.rs

@@ -82,6 +82,7 @@ pub(crate) struct RequestData {
     peer_addr: String,
     username: Option<String>,
     application_name: Option<String>,
+    user_agent: Option<String>,
     endpoint_id: Option<String>,
     database: Option<String>,
     project: Option<String>,
@@ -128,6 +129,7 @@ impl From<&RequestContextInner> for RequestData {
             timestamp: value.first_packet.naive_utc(),
             username: value.user.as_deref().map(String::from),
             application_name: value.application.as_deref().map(String::from),
+            user_agent: value.user_agent.as_deref().map(String::from),
             endpoint_id: value.endpoint_id.as_deref().map(String::from),
             database: value.dbname.as_deref().map(String::from),
             project: value.project.as_deref().map(String::from),
@@ -522,6 +524,7 @@ mod tests {
             .unwrap()
             .naive_utc(),
             application_name: Some("test".to_owned()),
+            user_agent: Some("test-user-agent".to_owned()),
             username: Some(hex::encode(rng.r#gen::<[u8; 4]>())),
             endpoint_id: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
             database: Some(hex::encode(rng.r#gen::<[u8; 16]>())),
@@ -610,15 +613,15 @@ mod tests {
         assert_eq!(
             file_stats,
             [
-                (1313105, 3, 6000),
-                (1313094, 3, 6000),
-                (1313153, 3, 6000),
-                (1313110, 3, 6000),
-                (1313246, 3, 6000),
-                (1313083, 3, 6000),
-                (1312877, 3, 6000),
-                (1313112, 3, 6000),
-                (438020, 1, 2000)
+                (1313953, 3, 6000),
+                (1313942, 3, 6000),
+                (1314001, 3, 6000),
+                (1313958, 3, 6000),
+                (1314094, 3, 6000),
+                (1313931, 3, 6000),
+                (1313725, 3, 6000),
+                (1313960, 3, 6000),
+                (438318, 1, 2000)
             ]
         );
 
@@ -650,11 +653,11 @@ mod tests {
         assert_eq!(
             file_stats,
             [
-                (1204324, 5, 10000),
-                (1204048, 5, 10000),
-                (1204349, 5, 10000),
-                (1204334, 5, 10000),
-                (1204588, 5, 10000)
+                (1205810, 5, 10000),
+                (1205534, 5, 10000),
+                (1205835, 5, 10000),
+                (1205820, 5, 10000),
+                (1206074, 5, 10000)
             ]
         );
 
@@ -679,15 +682,15 @@ mod tests {
         assert_eq!(
             file_stats,
             [
-                (1313105, 3, 6000),
-                (1313094, 3, 6000),
-                (1313153, 3, 6000),
-                (1313110, 3, 6000),
-                (1313246, 3, 6000),
-                (1313083, 3, 6000),
-                (1312877, 3, 6000),
-                (1313112, 3, 6000),
-                (438020, 1, 2000)
+                (1313953, 3, 6000),
+                (1313942, 3, 6000),
+                (1314001, 3, 6000),
+                (1313958, 3, 6000),
+                (1314094, 3, 6000),
+                (1313931, 3, 6000),
+                (1313725, 3, 6000),
+                (1313960, 3, 6000),
+                (438318, 1, 2000)
             ]
         );
 
@@ -724,7 +727,7 @@ mod tests {
         // files are smaller than the size threshold, but they took too long to fill so were flushed early
         assert_eq!(
             file_stats,
-            [(658014, 2, 3001), (657728, 2, 3000), (657524, 2, 2999)]
+            [(658584, 2, 3001), (658298, 2, 3000), (658094, 2, 2999)]
         );
 
         tmpdir.close().unwrap();

proxy/src/serverless/mod.rs

@@ -438,6 +438,14 @@ async fn request_handler(
             &config.region,
         );
 
+        ctx.set_user_agent(
+            request
+                .headers()
+                .get(hyper::header::USER_AGENT)
+                .and_then(|h| h.to_str().ok())
+                .map(Into::into),
+        );
+
         let span = ctx.span();
         info!(parent: &span, "performing websocket upgrade");
 

proxy/src/serverless/sql_over_http.rs

@@ -228,6 +228,13 @@ fn get_conn_info(
         }
     }
 
+    ctx.set_user_agent(
+        headers
+            .get(hyper::header::USER_AGENT)
+            .and_then(|h| h.to_str().ok())
+            .map(Into::into),
+    );
+
     let user_info = ComputeUserInfo {
         endpoint,
         user: username,