From c367c2c27418b90faa3ac561dc2697a4ac22be4e Mon Sep 17 00:00:00 2001
From: Egor Voynov <egor.voynov@aiven.io>
Date: Thu, 11 Apr 2024 16:27:17 +0200
Subject: [PATCH] pghoard: add xlog filename check to CLI commands

---
 golang/pghoard_postgres_command_go.go | 20 ++++++++++++--------
 pghoard/postgres_command.py           | 13 ++++++++-----
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/golang/pghoard_postgres_command_go.go b/golang/pghoard_postgres_command_go.go
index 1a6a9eba..e55161aa 100644
--- a/golang/pghoard_postgres_command_go.go
+++ b/golang/pghoard_postgres_command_go.go
@@ -16,6 +16,7 @@ import (
 	"net/http"
 	"os"
 	"path"
+	"regexp"
 	"time"
 )
 
@@ -120,15 +121,18 @@ func restore_command(url string, output string, xlog string) (int, error) {
 			}
 			output_path = path.Join(cwd, output)
 		}
-		// if file "<xlog>.pghoard.prefetch" exists, just move it to destination
-		xlogPrefetchPath := path.Join(path.Dir(output_path), xlog+".pghoard.prefetch")
-		_, err = os.Stat(xlogPrefetchPath)
-		if err == nil {
-			err := os.Rename(xlogPrefetchPath, output_path)
-			if err != nil {
-				return EXIT_ABORT, err
+		xlogNameRe := regexp.MustCompile(`^([A-F0-9]{24}|[A-F0-9]{8}\.history)$`)
+		if xlogNameRe.MatchString(xlog) {
+			// if file "<xlog>.pghoard.prefetch" exists, just move it to destination
+			xlogPrefetchPath := path.Join(path.Dir(output_path), xlog+".pghoard.prefetch")
+			_, err = os.Stat(xlogPrefetchPath)
+			if err == nil {
+				err := os.Rename(xlogPrefetchPath, output_path)
+				if err != nil {
+					return EXIT_ABORT, err
+				}
+				return EXIT_OK, nil
 			}
-			return EXIT_OK, nil
 		}
 		req, err = http.NewRequest("GET", url, nil)
 		req.Header.Set("x-pghoard-target-path", output_path)
diff --git a/pghoard/postgres_command.py b/pghoard/postgres_command.py
index 1067e4d1..df753ff5 100644
--- a/pghoard/postgres_command.py
+++ b/pghoard/postgres_command.py
@@ -12,6 +12,8 @@
 import time
 from http.client import BadStatusLine, HTTPConnection, IncompleteRead
 
+from pghoard.wal import TIMELINE_RE, WAL_RE
+
 from . import version
 
 PGHOARD_HOST = "127.0.0.1"
@@ -72,11 +74,12 @@ def restore_command(site, xlog, output, host=PGHOARD_HOST, port=PGHOARD_PORT, re
         # directory.  Note that os.path.join strips preceding components if a new components starts with a
         # slash so it's still possible to use this with absolute paths.
         output_path = os.path.join(os.getcwd(), output)
-        # if file "<xlog>.pghoard.prefetch" exists, just move it to destination
-        prefetch_path = os.path.join(os.path.dirname(output_path), xlog + ".pghoard.prefetch")
-        if os.path.exists(prefetch_path):
-            os.rename(prefetch_path, output_path)
-            return
+        if WAL_RE.match(xlog) or TIMELINE_RE.match(xlog):
+            # if file "<xlog>.pghoard.prefetch" exists, just move it to destination
+            prefetch_path = os.path.join(os.path.dirname(output_path), xlog + ".pghoard.prefetch")
+            if os.path.exists(prefetch_path):
+                os.rename(prefetch_path, output_path)
+                return
         headers = {"x-pghoard-target-path": output_path}
         method = "GET"
     path = "/{}/archive/{}".format(site, xlog)