From 49d1921a280b9ff681085bd629e5284a51bb4871 Mon Sep 17 00:00:00 2001
From: Patrick Insinger <patrickinsinger@gmail.com>
Date: Mon, 10 May 2021 17:16:40 -0400
Subject: [PATCH] page_server - add python api tests

---
 pageserver/src/page_service.rs     | 19 +++++++++++++++++++
 test_runner/test_pageserver_api.py | 19 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 test_runner/test_pageserver_api.py

diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index 3a77af46bb..52dab8856a 100644
--- a/pageserver/src/page_service.rs
+++ b/pageserver/src/page_service.rs
@@ -54,6 +54,7 @@ enum FeMessage {
 #[derive(Debug)]
 enum BeMessage {
     AuthenticationOk,
+    ParameterStatus,
     ReadyForQuery,
     RowDescription,
     ParseComplete,
@@ -439,6 +440,16 @@ impl Connection {
                 self.stream.write_i32::<BE>(0)?;
             }
 
+            BeMessage::ParameterStatus => {
+                self.stream.write_u8(b'S')?;
+                // parameter names and values are specified by null terminated strings
+                const PARAM_NAME_VALUE: &[u8] = b"client_encoding\0UTF8\0";
+                // length of this i32 + rest of data in message
+                self.stream
+                    .write_i32::<BE>(4 + PARAM_NAME_VALUE.len() as i32)?;
+                self.stream.write(PARAM_NAME_VALUE)?;
+            }
+
             BeMessage::ReadyForQuery => {
                 self.stream.write_u8(b'Z')?;
                 self.stream.write_i32::<BE>(4 + 1)?;
@@ -575,6 +586,9 @@ impl Connection {
                         }
                         StartupRequestCode::Normal => {
                             self.write_message_noflush(&BeMessage::AuthenticationOk)?;
+                            // psycopg2 will not connect if client_encoding is not
+                            // specified by the server
+                            self.write_message_noflush(&BeMessage::ParameterStatus)?;
                             self.write_message(&BeMessage::ReadyForQuery)?;
                             self.init_done = true;
                         }
@@ -676,6 +690,11 @@ impl Connection {
             self.write_message_noflush(&BeMessage::DataRow)?;
             self.write_message_noflush(&BeMessage::CommandComplete)?;
             self.write_message(&BeMessage::ReadyForQuery)?;
+        } else if query_string.to_ascii_lowercase().starts_with(b"set ") {
+            // important because psycopg2 executes "SET datestyle TO 'ISO'"
+            // on connect
+            self.write_message_noflush(&BeMessage::CommandComplete)?;
+            self.write_message(&BeMessage::ReadyForQuery)?;
         } else {
             self.write_message_noflush(&BeMessage::RowDescription)?;
             self.write_message_noflush(&BeMessage::DataRow)?;
diff --git a/test_runner/test_pageserver_api.py b/test_runner/test_pageserver_api.py
new file mode 100644
index 0000000000..b92e87001b
--- /dev/null
+++ b/test_runner/test_pageserver_api.py
@@ -0,0 +1,19 @@
+import pytest
+import psycopg2
+import getpass
+
+pytest_plugins = ("fixtures.zenith_fixtures")
+
+HOST = 'localhost'
+PAGESERVER_PORT = 64000
+
+def test_status(zen_simple):
+    username = getpass.getuser()
+    conn_str = 'host={} port={} dbname=postgres user={}'.format(
+        HOST, PAGESERVER_PORT, username)
+    pg_conn = psycopg2.connect(conn_str)
+    pg_conn.autocommit = True
+    cur = pg_conn.cursor()
+    cur.execute('status;')
+    assert cur.fetchone() == ('hello world',)
+    pg_conn.close()