[pgmoneta#15] Get extra files

GuChad369 · Aug 5, 2024 · 694a673 · 694a673
1 parent 33eb760
commit 694a673
Show file tree

Hide file tree

Showing 10 changed files with 668 additions and 351 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
diff --git a/doc/manual/user-02-functions.md b/doc/manual/user-02-functions.md
@@ -8,8 +8,10 @@ After you create the extension `pgmoneta_ext` using the `postgres` role, you can
 
 ## Extension functions
 
-| Function                    | Privilege                    | Description                                            |
-|-----------------------------|-----------|--------------------------------------------------------|
-| `pgmoneta_ext_version()`    |   Default                      | Return the version number of `pgmoneta_ext` as a Datum.|
-| `pgmoneta_ext_switch_wal()` | SUPERUSER                   | A function for switching to a new WAL file.            |
-| `pgmoneta_ext_checkpoint()` | SUPERUSER <br>pg_checkpoint | A function which forces a checkpoint. <br>This function can only be executed by a `SUPERUSER` in PostgreSQL 13/14, but can also be executed by `pg_checkpoint` in PostgreSQL 15+. <br>You can use the SQL command `GRANT pg_checkpoint TO repl;` to assign the role in PostgreSQL 15+.|
+| Function                    | Privilege | Parameters | Description                                            |
+|-----------------------------|-----------|------------|--------------------------------------------------------|
+| `pgmoneta_ext_version()`    |   Default |    None    | Return the version number of `pgmoneta_ext` as a Datum.|
+| `pgmoneta_ext_switch_wal()` | SUPERUSER |    None    | A function for switching to a new WAL file.            |
+| `pgmoneta_ext_checkpoint()` | SUPERUSER <br>pg_checkpoint | None | A function which forces a checkpoint. <br>This function can only be executed by a `SUPERUSER` in PostgreSQL 13/14, but can also be executed by `pg_checkpoint` in PostgreSQL 15+. <br>You can use the SQL command `GRANT pg_checkpoint TO repl;` to assign the role in PostgreSQL 15+.|
+| `pgmoneta_ext_get_file()`|   SUPERUSER        | path/to/file  | Return the bytes of the specified file that is passed in.|
+| `pgmoneta_ext_get_files()`    |   SUPERUSER    | path/to/dir   | Return all file paths in the specified directory passed in.|
diff --git a/sql/pgmoneta_ext--0.1.0.sql b/sql/pgmoneta_ext--0.1.0.sql
@@ -14,4 +14,12 @@ CREATE FUNCTION pgmoneta_ext_checkpoint(OUT success bool,
 ) 
 RETURNS record 
 AS 'MODULE_PATHNAME' 
-LANGUAGE C STRICT;
+LANGUAGE C STRICT;
+
+CREATE FUNCTION pgmoneta_ext_get_file(file_path TEXT) RETURNS TEXT
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT;
+
+CREATE OR REPLACE FUNCTION pgmoneta_ext_get_files(file_path TEXT) RETURNS text[]
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT;
diff --git a/src/include/utils.h b/src/include/utils.h
@@ -35,15 +35,21 @@ extern "C" {
 #endif
 
 /* PostgreSQL */
-#include "postgres.h"
-#include "fmgr.h"
+#include <postgres.h>
+#include <fmgr.h>
+#include <funcapi.h>
+#include <utils/array.h>
 
 /* system */
+#include <dirent.h>
 #include <stdbool.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
 
-#define PRIVILEDGE_DEFAULT            1 << 0  // 001
-#define PRIVILEDGE_PG_CHECKPOINT      1 << 1  // 010
-#define PRIVILEDGE_SUPERUSER          1 << 2  // 100
+#define PRIVILEGE_DEFAULT            1 << 0  // 001
+#define PRIVILEGE_PG_CHECKPOINT      1 << 1  // 010
+#define PRIVILEGE_SUPERUSER          1 << 2  // 100
 
 /**
  * Check if the role has superuser privileges.

diff --git a/src/pgmoneta_ext/lib.c b/src/pgmoneta_ext/lib.c
@@ -32,27 +32,50 @@
 #include <utils.h>
 
 /* PostgreSQL */
-#include "postgres.h"
-#include "access/htup_details.h"
-#include "access/xlog.h"
-#include "access/xlog_internal.h"
-#include "catalog/pg_authid.h"
-#include "catalog/pg_type.h"
-#include "fmgr.h"
-#include "funcapi.h"
-#include "miscadmin.h"
-#include "utils/acl.h"
-#include "utils/builtins.h"
-#include "utils/elog.h"
-#include "utils/errcodes.h"
-#include "utils/lsyscache.h"
-#include "utils/syscache.h"
+#include <postgres.h>
+#include <access/htup_details.h>
+#include <access/xlog.h>
+#include <access/xlog_internal.h>
+#include <catalog/pg_authid.h>
+#include <catalog/pg_type.h>
+#include <commands/dbcommands.h>
+#include <common/base64.h>
+#include <fmgr.h>
+#include <funcapi.h>
+#include <lib/stringinfo.h>
+#include <miscadmin.h>
+#include <storage/fd.h>
+#include <utils/acl.h>
+#include <utils/array.h>
+#include <utils/builtins.h>
+#include <utils/bytea.h>
+#include <utils/elog.h>
+#include <utils/errcodes.h>
+#include <utils/lsyscache.h>
+#include <utils/memutils.h>
+#include <mb/pg_wchar.h>
+#include <utils/syscache.h>
 
+/* system */
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#include <dirent.h>
+
+#ifdef PG_MODULE_MAGIC
 PG_MODULE_MAGIC;
+#endif
+
+#define PGMONETA_EXT_CHUNK_SIZE 8192
+
+static text* encode_bytea_to_base64(bytea* data);
+static void list_files(const char* name, ArrayBuildState* astate);
 
 PG_FUNCTION_INFO_V1(pgmoneta_ext_version);
 PG_FUNCTION_INFO_V1(pgmoneta_ext_switch_wal);
 PG_FUNCTION_INFO_V1(pgmoneta_ext_checkpoint);
+PG_FUNCTION_INFO_V1(pgmoneta_ext_get_file);
+PG_FUNCTION_INFO_V1(pgmoneta_ext_get_files);
 
 Datum
 pgmoneta_ext_version(PG_FUNCTION_ARGS)
@@ -87,7 +110,7 @@ pgmoneta_ext_switch_wal(PG_FUNCTION_ARGS)
 
    privileges = pgmoneta_ext_check_privilege(roleid);
 
-   if (privileges & PRIVILEDGE_SUPERUSER)
+   if (privileges & PRIVILEGE_SUPERUSER)
    {
       // Request to switch WAL with mark_unimportant set to false.
       recptr = RequestXLogSwitch(false);
@@ -140,7 +163,7 @@ pgmoneta_ext_checkpoint(PG_FUNCTION_ARGS)
    roleid = GetUserId();
    privileges = pgmoneta_ext_check_privilege(roleid);
 
-   if (privileges & (PRIVILEDGE_PG_CHECKPOINT | PRIVILEDGE_SUPERUSER))
+   if (privileges & (PRIVILEGE_PG_CHECKPOINT | PRIVILEGE_SUPERUSER))
    {
       // Perform the checkpoint
       RequestCheckpoint(CHECKPOINT_IMMEDIATE | CHECKPOINT_WAIT | CHECKPOINT_FORCE);
@@ -168,4 +191,179 @@ pgmoneta_ext_checkpoint(PG_FUNCTION_ARGS)
    result = HeapTupleGetDatum(tuple);
 
    PG_RETURN_DATUM(result);
+}
+
+Datum
+pgmoneta_ext_get_file(PG_FUNCTION_ARGS)
+{
+   int privileges;
+   char buffer[PGMONETA_EXT_CHUNK_SIZE];
+   char* file_path;
+   size_t bytes_read;
+   text* file_path_arg;
+   FILE* file;
+   StringInfoData result;
+   Oid roleid;
+   bytea* file_data;
+   text* base64_result;
+
+   file_path_arg = PG_GETARG_TEXT_PP(0);
+   file_path = text_to_cstring(file_path_arg);
+
+   roleid = GetUserId();
+   privileges = pgmoneta_ext_check_privilege(roleid);
+
+   if (privileges & PRIVILEGE_SUPERUSER)
+   {
+      initStringInfo(&result);
+
+      // Open the file
+      file = AllocateFile(file_path, "r");
+      if (file == NULL)
+      {
+         ereport(ERROR, errmsg_internal("pgmoneta_ext_get_file: Could not open file \"%s\": %m", file_path));
+         PG_RETURN_NULL();
+      }
+
+      PG_TRY();
+      {
+         // Read the file content in chunks
+         while ((bytes_read = fread(buffer, 1, sizeof(buffer), file)) > 0)
+         {
+            if (ferror(file))
+            {
+               ereport(ERROR, (errmsg_internal("pgmoneta_ext_get_file: Could not read file \"%s\": %m", file_path)));
+            }
+            appendBinaryStringInfo(&result, buffer, bytes_read);
+         }
+      }
+      PG_CATCH();
+      {
+         FreeFile(file);
+         PG_RE_THROW();
+      }
+      PG_END_TRY();
+
+      FreeFile(file);
+
+      file_data = (bytea*)palloc(VARHDRSZ + result.len);
+      SET_VARSIZE(file_data, VARHDRSZ + result.len);
+      memcpy(VARDATA(file_data), result.data, result.len);
+
+      // Encode as Base64
+      base64_result = encode_bytea_to_base64(file_data);
+
+      PG_RETURN_TEXT_P(base64_result);
+   }
+   else
+   {
+      ereport(LOG, errmsg_internal("pgmoneta_ext_checkpoint: Current role is not a superuser"));
+      PG_RETURN_NULL();
+   }
+}
+
+Datum
+pgmoneta_ext_get_files(PG_FUNCTION_ARGS)
+{
+   text* file_path_arg;
+   char* file_path;
+   struct stat path_stat;
+   ArrayBuildState* astate = NULL;
+   Oid roleid;
+   int privileges;
+
+   file_path_arg = PG_GETARG_TEXT_PP(0);
+   file_path = text_to_cstring(file_path_arg);
+
+   roleid = GetUserId();
+   privileges = pgmoneta_ext_check_privilege(roleid);
+
+   if (privileges & PRIVILEGE_SUPERUSER)
+   {
+      if (stat(file_path, &path_stat) != 0)
+      {
+         ereport(ERROR, errmsg_internal("Could not stat file or directory: %s", file_path));
+         PG_RETURN_NULL();
+      }
+
+      astate = initArrayResult(TEXTOID, CurrentMemoryContext, false);
+
+      if (S_ISDIR(path_stat.st_mode))
+      {
+         list_files(file_path, astate);
+      }
+      else if (S_ISREG(path_stat.st_mode))
+      {
+         astate = accumArrayResult(astate, CStringGetTextDatum(file_path), false, TEXTOID, CurrentMemoryContext);
+      }
+      else
+      {
+         ereport(ERROR, errmsg_internal("Argument is neither a file nor a directory: %s", file_path));
+         PG_RETURN_NULL();
+      }
+
+      // PG_RETURN_ARRAYTYPE_P(makeArrayResult(astate, CurrentMemoryContext));
+
+      PG_RETURN_ARRAYTYPE_P((ArrayType *) DatumGetPointer(makeArrayResult(astate, CurrentMemoryContext)));
+   }
+   else
+   {
+      ereport(LOG, errmsg_internal("pgmoneta_ext_checkpoint: Current role is not a superuser"));
+      PG_RETURN_NULL();
+   }
+}
+
+static text*
+encode_bytea_to_base64(bytea* data)
+{
+   size_t data_len = VARSIZE_ANY_EXHDR(data);
+   size_t encoded_len = pg_b64_enc_len(data_len);
+   char* encoded = palloc(encoded_len + 1);
+   int actual_encoded_len = pg_b64_encode(VARDATA_ANY(data), data_len, encoded, encoded_len);
+
+   text* result = (text*)palloc(VARHDRSZ + actual_encoded_len);
+   SET_VARSIZE(result, VARHDRSZ + actual_encoded_len);
+   memcpy(VARDATA(result), encoded, actual_encoded_len);
+
+   pfree(encoded);
+   return result;
+}
+
+static void
+list_files(const char* name, ArrayBuildState* astate)
+{
+   DIR* dir;
+   struct dirent* entry;
+
+   if (!(dir = opendir(name)))
+   {
+      return;
+   }
+
+   while ((entry = readdir(dir)) != NULL)
+   {
+      char path[PATH_MAX];
+      snprintf(path, sizeof(path), "%s/%s", name, entry->d_name);
+
+      struct stat path_stat;
+      if (stat(path, &path_stat) != 0)
+      {
+         continue;
+      }
+
+      if (S_ISDIR(path_stat.st_mode))
+      {
+         if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0)
+         {
+            continue;
+         }
+
+         list_files(path, astate);
+      }
+      else if (S_ISREG(path_stat.st_mode))
+      {
+         astate = accumArrayResult(astate, CStringGetTextDatum(path), false, TEXTOID, CurrentMemoryContext);
+      }
+   }
+   closedir(dir);
 }
diff --git a/src/pgmoneta_ext/utils.c b/src/pgmoneta_ext/utils.c
@@ -31,15 +31,15 @@
 #include <utils.h>
 
 /* PostgreSQL */
-#include "access/htup_details.h"
-#include "catalog/pg_authid.h"
-#include "miscadmin.h"
-#include "utils/acl.h"
-#include "utils/builtins.h"
-#include "utils/elog.h"
-#include "utils/errcodes.h"
-#include "utils/relcache.h"
-#include "utils/syscache.h"
+#include <access/htup_details.h>
+#include <catalog/pg_authid.h>
+#include <miscadmin.h>
+#include <utils/acl.h>
+#include <utils/builtins.h>
+#include <utils/elog.h>
+#include <utils/errcodes.h>
+#include <utils/relcache.h>
+#include <utils/syscache.h>
 
 /* system */
 
@@ -116,19 +116,19 @@ pgmoneta_ext_check_privilege(Oid roleid)
 
    if (pgmoneta_ext_check_superuser(roleid))
    {
-      privileges |= PRIVILEDGE_SUPERUSER;
+      privileges |= PRIVILEGE_SUPERUSER;
    }
 
 #if PG_VERSION_NUM >= 150000
    bool is_pg_checkpoint = pgmoneta_ext_check_role(roleid, "pg_checkpoint");
 
    if (is_pg_checkpoint)
    {
-      privileges |= PRIVILEDGE_PG_CHECKPOINT;
+      privileges |= PRIVILEGE_PG_CHECKPOINT;
    }
 #endif
 
-   privileges |= PRIVILEDGE_DEFAULT;
+   privileges |= PRIVILEGE_DEFAULT;
 
    return privileges;
-}
+}
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
@@ -51,6 +51,7 @@ if(container)
   add_test(test_version_14_rocky9 "${CMAKE_CURRENT_SOURCE_DIR}/../test/testsuite.sh" "${CMAKE_CURRENT_SOURCE_DIR}/../test" "Dockerfile.rocky9" 14)
   add_test(test_version_15_rocky9 "${CMAKE_CURRENT_SOURCE_DIR}/../test/testsuite.sh" "${CMAKE_CURRENT_SOURCE_DIR}/../test" "Dockerfile.rocky9" 15)
   add_test(test_version_16_rocky9 "${CMAKE_CURRENT_SOURCE_DIR}/../test/testsuite.sh" "${CMAKE_CURRENT_SOURCE_DIR}/../test" "Dockerfile.rocky9" 16)
+  add_test(test_version_17_rocky9 "${CMAKE_CURRENT_SOURCE_DIR}/../test/testsuite.sh" "${CMAKE_CURRENT_SOURCE_DIR}/../test" "Dockerfile.rocky9" 17)
 
   set(failRegex 
       "Failures: [1-9][0-9]*"
@@ -64,4 +65,6 @@ if(container)
                 PROPERTY FAIL_REGULAR_EXPRESSION "${failRegex}")        
   set_property (TEST test_version_16_rocky9
                 PROPERTY FAIL_REGULAR_EXPRESSION "${failRegex}")
+  set_property (TEST test_version_17_rocky9
+                PROPERTY FAIL_REGULAR_EXPRESSION "${failRegex}")
 endif()