From aafdd9c454269324982e41a72d6973582b49132b Mon Sep 17 00:00:00 2001
From: Jubilee101 <haoranz4@andrew.cmu.edu>
Date: Fri, 26 May 2023 19:11:29 -0400
Subject: [PATCH] [#123]Common functionalities for querying server version

---
 src/include/message.h     |  9 +++++++++
 src/include/pgmoneta.h    |  1 +
 src/include/server.h      |  9 +++++++++
 src/libpgmoneta/message.c | 27 +++++++++++++++++++++++++++
 src/libpgmoneta/server.c  | 32 ++++++++++++++++++++++++++++++++
 5 files changed, 78 insertions(+)

diff --git a/src/include/message.h b/src/include/message.h
index 14b1891d..1d3969f5 100644
--- a/src/include/message.h
+++ b/src/include/message.h
@@ -331,6 +331,15 @@ pgmoneta_create_start_replication_message(char* xlogpos, int timeline, char* slo
 int
 pgmoneta_create_standby_status_update_message(int64_t received, int64_t flushed, int64_t applied, struct message** msg);
 
+/**
+ * Create a query message for a simple query
+ * @param query The query to be executed on server
+ * @param msg The resulting message
+ * @return 0 upon success, otherwise 1
+ */
+int
+pgmoneta_create_query_message(char* query, struct message** msg);
+
 /**
  * Has a message
  * @param type The message type to be extracted
diff --git a/src/include/pgmoneta.h b/src/include/pgmoneta.h
index a3dab0e6..447138e5 100644
--- a/src/include/pgmoneta.h
+++ b/src/include/pgmoneta.h
@@ -193,6 +193,7 @@ struct server
    int wal_size;                       /**< The size of the WAL files */
    bool wal_streaming;                 /**< Is WAL streaming active */
    bool valid;                         /**< Is the server valid */
+   int version;                        /**< The major version of the server*/
 } __attribute__ ((aligned (64)));
 
 /** @struct
diff --git a/src/include/server.h b/src/include/server.h
index 1a751bf1..016c1097 100644
--- a/src/include/server.h
+++ b/src/include/server.h
@@ -51,6 +51,15 @@ pgmoneta_server_info(int srv);
 int
 pgmoneta_server_get_wal_size(int socket, int* ws);
 
+/**
+ * Get the version of a server
+ * @param socket The socket
+ * @param server The server index
+ * @return 0 upon success, otherwise 1
+ */
+int
+pgmoneta_server_get_version(int socket, int server);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/libpgmoneta/message.c b/src/libpgmoneta/message.c
index 43df3c49..c1b5f70c 100644
--- a/src/libpgmoneta/message.c
+++ b/src/libpgmoneta/message.c
@@ -735,6 +735,33 @@ pgmoneta_create_standby_status_update_message(int64_t received, int64_t flushed,
    return MESSAGE_STATUS_OK;
 }
 
+int
+pgmoneta_create_query_message(char* query, struct message** msg)
+{
+   struct message* m = NULL;
+   size_t size;
+   char cmd[1024];
+
+   memset(&cmd[0], 0, sizeof(cmd));
+   strcpy(cmd, query);
+   size = 1 + 4 + strlen(cmd) + 1;
+   m = (struct message*)malloc(sizeof(struct message));
+   m->data = malloc(size);
+
+   memset(m->data, 0, size);
+
+   m->kind = 'Q';
+   m->length = size;
+
+   pgmoneta_write_byte(m->data, 'Q');
+   pgmoneta_write_int32(m->data + 1, size - 1);
+   memcpy(m->data + 5, &cmd[0], strlen(cmd));
+
+   *msg = m;
+
+   return MESSAGE_STATUS_OK;
+}
+
 int
 pgmoneta_query_execute(int socket, struct message* msg, struct query_response** response)
 {
diff --git a/src/libpgmoneta/server.c b/src/libpgmoneta/server.c
index 8920d2b0..d93a9b29 100644
--- a/src/libpgmoneta/server.c
+++ b/src/libpgmoneta/server.c
@@ -257,3 +257,35 @@ get_wal_level(int socket, bool* replica)
 
    return 1;
 }
+
+int
+pgmoneta_server_get_version(int socket, int server)
+{
+   int ret;
+   struct message* query_msg = NULL;
+   struct query_response* response = NULL;
+   struct configuration* config;
+
+   config = (struct configuration*)shmem;
+
+   ret = pgmoneta_create_query_message("SELECT split_part(split_part(version(), ' ', 2), '.', 1);", &query_msg);
+   if (ret != MESSAGE_STATUS_OK)
+   {
+      goto error;
+   }
+
+   if (pgmoneta_query_execute(socket, query_msg, &response) || response == NULL)
+   {
+      goto error;
+   }
+
+   config->servers[server].version = atoi(response->tuples->data[0]);
+
+   pgmoneta_free_query_response(response);
+   pgmoneta_free_message(query_msg);
+
+   return 0;
+error:
+   pgmoneta_free_query_response(response);
+   return 1;
+}