From 3b2159310f7026686c29b27d2124bfceec58e8df Mon Sep 17 00:00:00 2001
From: Andrey Rudenko <me@prepor.dev>
Date: Thu, 13 Jun 2024 10:52:41 +0200
Subject: [PATCH] fixes to protocol based on
 https://github.com/neondatabase/neon/pull/8039

---
 compute_tools/src/compute.rs   | 19 ++++++++++++++++---
 pageserver/src/page_service.rs |  5 ++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index 54e51d0557..b494c92722 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -14,6 +14,7 @@ use std::time;
 use std::time::Instant;
 use std::time::SystemTime;
 
+use anyhow::bail;
 use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
 use futures::future::join_all;
@@ -21,6 +22,7 @@ use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use nix::unistd::Pid;
 use postgres::error::SqlState;
+use postgres::SimpleQueryMessage;
 use postgres::{Client, NoTls};
 use tracing::{debug, error, info, instrument, warn};
 use utils::id::{TenantId, TimelineId};
@@ -1452,9 +1454,20 @@ fn lsn_lease_request(configs: &[postgres::Config], cmd: &str) -> Result<u128> {
         .iter()
         .map(|config| {
             let mut client = config.connect(NoTls)?;
-            let msg = client.query_one(cmd, &[])?;
-            let bytes: &[u8] = msg.try_get("valid_until")?;
-            Ok(u128::from_be_bytes(bytes.try_into()?))
+            let res = client.simple_query(cmd)?;
+            let msg = match res.get(0) {
+                Some(msg) => msg,
+                None => bail!("empty response"),
+            };
+            let row = match msg {
+                SimpleQueryMessage::Row(row) => row,
+                _ => bail!("error parsing lsn lease response"),
+            };
+            let valid_until_str = match row.get("valid_until") {
+                Some(valid_until) => valid_until,
+                None => bail!("valid_until not found"),
+            };
+            Ok(u128::from_str(valid_until_str)?)
         })
         .collect::<Result<Vec<u128>>>()?
         .into_iter()
diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index ae389826d5..35021fe6f1 100644
--- a/pageserver/src/page_service.rs
+++ b/pageserver/src/page_service.rs
@@ -945,9 +945,8 @@ impl PageServerHandler {
             b"valid_until",
         )]))?
         .write_message_noflush(&BeMessage::DataRow(&[Some(
-            &valid_until.as_millis().to_be_bytes(),
-        )]))?
-        .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+            valid_until.as_millis().to_string().as_bytes(),
+        )]))?;
 
         Ok(())
     }