File: barman-wal-restore

package info (click to toggle)
barman-cli 1.2-1
  • links: PTS, VCS
  • area: main
  • in suites: stretch
  • size: 160 kB
  • sloc: python: 292; makefile: 12
file content (430 lines) | stat: -rwxr-xr-x 14,357 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
#!/usr/bin/env python
#
# barman-wal-restore - Remote Barman WAL restore command for PostgreSQL
#
# This script remotely fetches WAL files from Barman via SSH, on demand.
# It is intended to be used as restore_command in recovery.conf files
# of PostgreSQL standby servers. Supports parallel fetching and
# protects against SSH failures.
#
# See the help page for usage information.
#
# Copyright (C) 2016 2ndQuadrant Italia Srl <info@2ndquadrant.it>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function

import os
import shutil
import subprocess
import sys
import time

try:
    import argparse
except ImportError:
    raise SystemExit("Missing required python module: argparse")

DEFAULT_USER = 'barman'
# TODO: Remove version from here and use setup.py's
VERSION = '1.2'
# TODO: make this generic
SPOOL_DIR = '/var/tmp/barman-wal-restore'

# The string_types list is used to identify strings
# in a consistent way between python 2 and 3
if sys.version_info[0] == 3:
    string_types = str,
else:
    string_types = basestring,


def main():
    """
    The main script entry point
    """
    config = parse_arguments()

    # Check WAL destination is not a directory
    if os.path.isdir(config.wal_dest):
        exit_with_error("WAL_DEST cannot be a directory: %s" %
                        config.wal_dest)

    # Open the destination file
    try:
        dest_file = open(config.wal_dest, 'wb')
    except EnvironmentError as e:
        exit_with_error("Cannot open '%s' (WAL_DEST) for writing: %s" %
                        (config.wal_dest, e))

    # If the file is present in SPOOL_DIR use it and terminate
    try_deliver_from_spool(config, dest_file)

    # If required load the list of files to download in parallel
    additional_files = peek_additional_files(config)

    try:
        # Execute barman get-wal through the ssh connection
        ssh_process = RemoteGetWal(config, config.wal_name, dest_file)
    except EnvironmentError as e:
        exit_with_error('Error executing "ssh": %s' % e, sleep=config.sleep)

    # Spawn a process for every additional file
    parallel_ssh_processes = spawn_additional_process(
        config, additional_files)

    # Wait for termination of every subprocess. If CTRL+C is pressed,
    # terminate all of them
    try:
        RemoteGetWal.wait_for_all()
    finally:
        # Cleanup failed spool files in case of errors
        for process in parallel_ssh_processes:
            if process.returncode != 0:
                os.unlink(process.dest_file)

    # If the command succeeded exit here
    if ssh_process.returncode == 0:
        sys.exit(0)

    # Report the exit code, remapping ssh failure code (255) to 3
    if ssh_process.returncode == 255:
        exit_with_error("Connection problem with ssh", 3, sleep=config.sleep)
    else:
        exit_with_error("Remote 'barman get-wal' command has failed!",
                        ssh_process.returncode, sleep=config.sleep)


def spawn_additional_process(config, additional_files):
    """
    Execute additional barman get-wal processes

    :param argparse.Namespace config: the configuration from command line
    :param additional_files: A list of WAL file to be downloaded in parallel
    :return list[subprocess.Popen]: list of created processes
    """
    processes = []
    for wal_name in additional_files:
        spool_file_name = os.path.join(SPOOL_DIR, wal_name)
        try:
            # Spawn a process and write the output in the spool dir
            process = RemoteGetWal(config, wal_name, spool_file_name)
            processes.append(process)
        except EnvironmentError:
            # If execution has failed make sure the spool file is unlinked
            try:
                os.unlink(spool_file_name)
            except EnvironmentError:
                # Suppress unlink errors
                pass

    return processes


def peek_additional_files(config):
    """
    Invoke remote get-wal --peek to receive a list of wal files to copy

    :param argparse.Namespace config: the configuration from command line
    :returns set: a set of WAL file names from the peek command
    """
    # If parallel downloading is not required return an empty array
    if not config.parallel:
        return []

    # Make sure the SPOOL_DIR exists
    try:
        if not os.path.exists(SPOOL_DIR):
            os.mkdir(SPOOL_DIR)
    except EnvironmentError as e:
        exit_with_error("Cannot create '%s' directory: %s" % (SPOOL_DIR, e))

    # Retrieve the list of files from remote
    additional_files = execute_peek(config)

    # Sanity check
    if len(additional_files) == 0 or additional_files[0] != config.wal_name:
        exit_with_error("The required file is not available: %s" %
                        config.wal_name)

    # Remove the first element, as now we know is identical to config.wal_name
    del additional_files[0]

    return additional_files


def build_ssh_command(config, wal_name, peek=0):
    """
    Prepare an ssh command according to the arguments passed on command line

    :param argparse.Namespace config: the configuration from command line
    :param str wal_name: the wal_name get-wal parameter
    :param int peek: in
    :return list[str]: the ssh command as list of string
    """
    ssh_command = [
        'ssh',
        "%s@%s" % (config.user, config.barman_host),
        "barman",
    ]

    if config.config:
        ssh_command.append("--config %s" % config.config)

    if peek:
        get_wal_command = "get-wal --peek '%s' '%s' '%s'" % (
            peek, config.server_name, wal_name)

    else:
        if config.compression:
            get_wal_command = "get-wal --%s '%s' '%s'" % (
                config.compression, config.server_name, wal_name)
        else:
            get_wal_command = "get-wal '%s' '%s'" % (
                config.server_name, wal_name)

    ssh_command.append(get_wal_command)
    return ssh_command


def execute_peek(config):
    """
    Invoke remote get-wall --peek to receive a list of wal file to copy

    :param argparse.Namespace config: the configuration from command line
    :returns set: a set of WAL file names from the peek command
    """
    # Build the peek command
    ssh_command = build_ssh_command(config, config.wal_name, config.parallel)
    # Issue the command
    try:
        output = subprocess.Popen(ssh_command,
                                  stdout=subprocess.PIPE).communicate()
        return list(output[0].decode().splitlines())
    except subprocess.CalledProcessError as e:
        exit_with_error("Impossible to invoke remote get-wall --peek: %s" % e)


def try_deliver_from_spool(config, dest_file):
    """
    Search for the requested file in the spool directory.
    If is already present, then copy it locally and exit,
    return otherwise.

    :param argparse.Namespace config: the configuration from command line
    :param dest_file: The destination file object
    """
    spool_file = os.path.join(SPOOL_DIR, config.wal_name)

    # id the file is not present, give up
    if not os.path.exists(spool_file):
        return

    try:
        shutil.copyfileobj(open(spool_file, 'rb'), dest_file)
        os.unlink(spool_file)
        sys.exit(0)
    except IOError as e:
        exit_with_error("Failure copying %s to %s: %s" %
                        (spool_file, dest_file.name, e))


def exit_with_error(message, status=2, sleep=0):
    """
    Print ``message`` and terminate the script with ``status``

    :param str message: message to print
    :param int status: script exit code
    :param int sleep: second to sleep before exiting
    """
    print("ERROR: %s" % message, file=sys.stderr)
    # Sleep for config.sleep seconds if required
    if sleep:
        print("Sleeping for %d seconds." % sleep, file=sys.stderr)
        time.sleep(sleep)
    sys.exit(status)


def parse_arguments(args=None):
    """
    Parse the command line arguments

    :param list[str] args: the raw arguments list. When not provided
        it defaults to sys.args[1:]
    :rtype: argparse.Namespace
    """
    parser = argparse.ArgumentParser(
        description="This script will be used as a 'restore_command' "
                    "based on the get-wal feature of Barman. "
                    "A ssh connection will be opened to the Barman host.",
    )
    parser.add_argument(
        '-V', '--version',
        action='version', version='%%(prog)s %s' % VERSION)
    parser.add_argument(
        "-U", "--user", default=DEFAULT_USER,
        help="The user used for the ssh connection to the Barman server. "
             "Defaults to '%(default)s'.",
    )
    parser.add_argument(
        "-s", "--sleep", default=0,
        type=int,
        metavar="SECONDS",
        help="Sleep for SECONDS after a failure of get-wal request. "
             "Defaults to 0 (nowait).",
    )
    parser.add_argument(
        "-p", "--parallel", default=0,
        type=int,
        metavar="JOBS",
        help="Specifies the number of files to peek and transfer "
             "in parallel. "
             "Defaults to 0 (disabled).",
    )
    parser.add_argument(
        '-z', '--gzip',
        help='Transfer the WAL files compressed with gzip',
        action='store_const', const='gzip', dest='compression',
    )
    parser.add_argument(
        '-j', '--bzip2',
        help='Transfer the WAL files compressed with bzip2',
        action='store_const', const='bzip2', dest='compression',
    )
    parser.add_argument(
        '-c', '--config',
        metavar="CONFIG",
        help='configuration file on the Barman server',
    )
    parser.add_argument(
        "barman_host",
        metavar="BARMAN_HOST",
        help="The host of the Barman server.",
    )
    parser.add_argument(
        "server_name",
        metavar="SERVER_NAME",
        help="The server name configured in Barman "
             "from which WALs are taken.",
    )
    parser.add_argument(
        "wal_name",
        metavar="WAL_NAME",
        help="The value of the '%%f' keyword "
             "(according to 'restore_command').",
    )
    parser.add_argument(
        "wal_dest",
        metavar="WAL_DEST",
        help="The value of the '%%p' keyword "
             "(according to 'restore_command').",
    )
    return parser.parse_args(args=args)


class RemoteGetWal(object):

    processes = set()
    """
    The list of processes that has been spawned by RemoteGetWal
    """

    def __init__(self, config, wal_name, dest_file):
        """
        Spawn a process that download a WAL from remote.

        If needed decompress the remote stream on the fly.

        :param argparse.Namespace config: the configuration from command line
        :param wal_name: The name of WAL to download
        :param dest_file: The destination file name or a writable file object
        """
        self.config = config
        self.wal_name = wal_name
        self.decompressor = None
        self.dest_file = None

        # If a string has been passed, it's the name of the destination file
        # We convert it in a writable binary file object
        if isinstance(dest_file, string_types):
            self.dest_file = dest_file
            dest_file = open(dest_file, 'wb')

        with dest_file:
            # If compression has been required, we need to spawn two processes
            if config.compression:
                # Spawn a remote get-wal process
                self.ssh_process = subprocess.Popen(
                    build_ssh_command(config, wal_name),
                    stdout=subprocess.PIPE)
                # Spawn the local decompressor
                self.decompressor = subprocess.Popen(
                    [config.compression, '-d'],
                    stdin=self.ssh_process.stdout,
                    stdout=dest_file
                )
                # Close the pipe descriptor, letting the decompressor process
                # to receive the SIGPIPE
                self.ssh_process.stdout.close()
            else:
                # With no compression only the remote get-wal process
                # is required
                self.ssh_process = subprocess.Popen(
                    build_ssh_command(config, wal_name),
                    stdout=dest_file)

        # Register the spawned processes in the class registry
        self.processes.add(self.ssh_process)
        if self.decompressor:
            self.processes.add(self.decompressor)

    @classmethod
    def wait_for_all(cls):
        """
        Wait for the termination of all the registered spawned processes.
        """
        try:
            while len(cls.processes):
                time.sleep(0.1)
                for process in cls.processes.copy():
                    if process.poll() is not None:
                        cls.processes.remove(process)
        except KeyboardInterrupt:
            # If a SIGINT has been received, make sure that every subprocess
            # terminate
            for process in cls.processes:
                process.kill()
            exit_with_error('SIGINT received! Terminating.')

    @property
    def returncode(self):
        """
        Return the exit code of the RemoteGetWal processes.

        A remote get-wal process return code is 0 only if both the remote
        get-wal process and the eventual decompressor return 0

        :return: exit code of the RemoteGetWal processes
        """
        if self.ssh_process.returncode != 0:
            return self.ssh_process.returncode
        if self.decompressor:
            if self.decompressor.returncode != 0:
                return self.decompressor.returncode
        return 0


if __name__ == '__main__':
    main()