[Issue #79] added tests for partial restore

tests/helpers/ptrack_helpers.py

@@ -130,7 +130,7 @@ def slow_start(self, replica=False):
     self.start()
     while True:
         try:
-            if self.safe_psql('postgres', query) == 't\n':
+            if self.safe_psql('template1', query) == 't\n':
                 break
         except testgres.QueryException as e:
             if 'database system is starting up' in e[0]:
@@ -296,8 +296,6 @@ class ProbackupTest(object):
 #                print('PGPROBACKUP_SSH_USER is not set')
 #                exit(1)
 
-
-
     def make_simple_node(
             self,
             base_dir=None,
@@ -342,6 +340,10 @@ class ProbackupTest(object):
                 'postgresql.auto.conf',
                 'max_wal_senders = 10')
 
+        # set major version
+        with open(os.path.join(node.data_dir, 'PG_VERSION')) as f:
+            node.major_version = f.read().rstrip()
+
         return node
 
     def create_tblspace_in_node(self, node, tblspc_name, tblspc_path=None, cfs=False):
@@ -584,6 +586,13 @@ class ProbackupTest(object):
 
         return filelist_diff
 
+    # used for partial restore
+    def truncate_every_file_in_dir(self, path):
+        for file in os.listdir(path):
+            print(file)
+            with open(os.path.join(path, file), "w") as f:
+                f.close()
+
     def check_ptrack_recovery(self, idx_dict):
         size = idx_dict['size']
         for PageNum in range(size):
@@ -1309,8 +1318,6 @@ class ProbackupTest(object):
                         os.path.join(restored_pgdata['pgdata'], directory),
                         restored_pgdata['dirs'][directory]['mode'])
 
-
-
         for directory in original_pgdata['dirs']:
             if directory not in restored_pgdata['dirs']:
                 fail = True
@@ -1318,7 +1325,6 @@ class ProbackupTest(object):
                 error_message += ' in restored PGDATA: {0}\n'.format(
                     os.path.join(restored_pgdata['pgdata'], directory))
 
-
         for file in restored_pgdata['files']:
             # File is present in RESTORED PGDATA
             # but not present in ORIGINAL

tests/restore.py

@@ -8,6 +8,8 @@ from time import sleep
 from datetime import datetime, timedelta
 import hashlib
 import shutil
+import json
+from testgres import QueryException
 
 
 module_name = 'restore'
@@ -2340,3 +2342,421 @@ class RestoreTest(ProbackupTest, unittest.TestCase):
 
         # Clean after yourself
         self.del_test_dir(module_name, fname)
+
+    @unittest.skip("skip")
+    def test_restore_specific_database_proof_of_concept(self):
+        """"""
+        fname = self.id().split('.')[3]
+        backup_dir = os.path.join(self.tmp_path, module_name, fname, 'backup')
+        node = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node'),
+            initdb_params=['--data-checksums'],
+            pg_options={
+                'autovacuum': 'off',
+                'shared_buffers': '512MB'})
+
+        self.init_pb(backup_dir)
+        self.add_instance(backup_dir, 'node', node)
+        self.set_archiving(backup_dir, 'node', node)
+        node.slow_start()
+
+        node.pgbench_init(scale=200)
+
+        exit(1)
+
+        # FULL backup
+        backup_id = self.backup_node(backup_dir, 'node', node)
+
+        pgbench = node.pgbench(
+            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+            options=['-T', '100', '-c8', '-j2', '--no-vacuum'])
+
+        pgbench.wait()
+        pgbench.stdout.close()
+
+        node_restored = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node_restored'))
+        node_restored.cleanup()
+
+        self.restore_node(backup_dir, 'node', node_restored)
+
+        node_restored.append_conf(
+            "postgresql.auto.conf", "port = {0}".format(node_restored.port))
+
+        # Clean after yourself
+        self.del_test_dir(module_name, fname)
+
+    def test_partial_restore_exclude(self):
+        """"""
+        fname = self.id().split('.')[3]
+        backup_dir = os.path.join(self.tmp_path, module_name, fname, 'backup')
+        node = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node'),
+            initdb_params=['--data-checksums'])
+
+        self.init_pb(backup_dir)
+        self.add_instance(backup_dir, 'node', node)
+        self.set_archiving(backup_dir, 'node', node)
+        node.slow_start()
+
+        for i in range(1, 10, 1):
+            node.safe_psql(
+                'postgres',
+                'CREATE database db{0}'.format(i))
+
+        db_list_raw = node.safe_psql(
+            'postgres',
+            'SELECT to_json(a) '
+            'FROM (SELECT oid, datname FROM pg_database) a').rstrip()
+
+        db_list_splitted = db_list_raw.splitlines()
+
+        db_list = {}
+        for line in db_list_splitted:
+            line = json.loads(line)
+            db_list[line['datname']] = line['oid']
+
+        # FULL backup
+        backup_id = self.backup_node(backup_dir, 'node', node)
+        pgdata = self.pgdata_content(node.data_dir)
+
+        # restore FULL backup
+        node_restored_1 = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node_restored_1'))
+        node_restored_1.cleanup()
+
+        try:
+            self.restore_node(
+                backup_dir, 'node',
+                node_restored_1, options=[
+                    "--db-include='db1'",
+                    "--db-exclude='db2'"])
+            self.assertEqual(
+                1, 0,
+                "Expecting Error because of 'db-exclude' and 'db-include'.\n "
+                "Output: {0} \n CMD: {1}".format(
+                    self.output, self.cmd))
+        except ProbackupException as e:
+            self.assertIn(
+                "ERROR: You cannot specify '--db-include' "
+                "and '--db-exclude' together", e.message,
+                '\n Unexpected Error Message: {0}\n CMD: {1}'.format(
+                    repr(e.message), self.cmd))
+
+        self.restore_node(
+            backup_dir, 'node', node_restored_1)
+
+        pgdata_restored_1 = self.pgdata_content(node_restored_1.data_dir)
+        self.compare_pgdata(pgdata, pgdata_restored_1)
+
+        db1_path = os.path.join(
+            node_restored_1.data_dir, 'base', db_list['db1'])
+        db5_path = os.path.join(
+            node_restored_1.data_dir, 'base', db_list['db5'])
+
+        self.truncate_every_file_in_dir(db1_path)
+        self.truncate_every_file_in_dir(db5_path)
+        pgdata_restored_1 = self.pgdata_content(node_restored_1.data_dir)
+
+        node_restored_2 = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node_restored_2'))
+        node_restored_2.cleanup()
+
+        self.restore_node(
+            backup_dir, 'node',
+            node_restored_2, options=[
+                "--db-exclude=db1",
+                "--db-exclude=db5"])
+
+        pgdata_restored_2 = self.pgdata_content(node_restored_2.data_dir)
+        self.compare_pgdata(pgdata_restored_1, pgdata_restored_2)
+
+        node_restored_2.append_conf(
+            "postgresql.auto.conf", "port = {0}".format(node_restored_2.port))
+
+        node_restored_2.slow_start()
+
+        node_restored_2.safe_psql(
+            'postgres',
+            'select 1')
+
+        try:
+            node_restored_2.safe_psql(
+                'db1',
+                'select 1')
+        except QueryException as e:
+            self.assertIn('FATAL', e.message)
+
+        try:
+            node_restored_2.safe_psql(
+                'db5',
+                'select 1')
+        except QueryException as e:
+            self.assertIn('FATAL', e.message)
+
+        with open(node_restored_2.pg_log_file, 'r') as f:
+            output = f.read()
+
+        self.assertNotIn('PANIC', output)
+
+        # Clean after yourself
+        self.del_test_dir(module_name, fname)
+
+    def test_partial_restore_exclude_tablespace(self):
+        """"""
+        fname = self.id().split('.')[3]
+        backup_dir = os.path.join(self.tmp_path, module_name, fname, 'backup')
+        node = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node'),
+            initdb_params=['--data-checksums'])
+
+        self.init_pb(backup_dir)
+        self.add_instance(backup_dir, 'node', node)
+        self.set_archiving(backup_dir, 'node', node)
+        node.slow_start()
+
+        cat_version = node.get_control_data()["Catalog version number"]
+        version_specific_dir = 'PG_' + node.major_version + '_' + cat_version
+
+        # PG_10_201707211
+        # pg_tblspc/33172/PG_9.5_201510051/16386/
+
+        self.create_tblspace_in_node(node, 'somedata')
+
+        node_tablespace = self.get_tblspace_path(node, 'somedata')
+
+        tbl_oid = node.safe_psql(
+            'postgres',
+            "SELECT oid "
+            "FROM pg_tablespace "
+            "WHERE spcname = 'somedata'").rstrip()
+
+        for i in range(1, 10, 1):
+            node.safe_psql(
+                'postgres',
+                'CREATE database db{0} tablespace somedata'.format(i))
+
+        db_list_raw = node.safe_psql(
+            'postgres',
+            'SELECT to_json(a) '
+            'FROM (SELECT oid, datname FROM pg_database) a').rstrip()
+
+        db_list_splitted = db_list_raw.splitlines()
+
+        db_list = {}
+        for line in db_list_splitted:
+            line = json.loads(line)
+            db_list[line['datname']] = line['oid']
+
+        # FULL backup
+        backup_id = self.backup_node(backup_dir, 'node', node)
+        pgdata = self.pgdata_content(node.data_dir)
+
+        # restore FULL backup
+        node_restored_1 = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node_restored_1'))
+        node_restored_1.cleanup()
+
+        node1_tablespace = self.get_tblspace_path(node_restored_1, 'somedata')
+
+        self.restore_node(
+            backup_dir, 'node',
+            node_restored_1, options=[
+                "-T", "{0}={1}".format(
+                    node_tablespace, node1_tablespace)])
+
+        pgdata_restored_1 = self.pgdata_content(node_restored_1.data_dir)
+        self.compare_pgdata(pgdata, pgdata_restored_1)
+
+        # truncate every db
+        for db in db_list:
+            # with exception below
+            if db in ['db1', 'db5']:
+                self.truncate_every_file_in_dir(
+                    os.path.join(
+                        node_restored_1.data_dir, 'pg_tblspc',
+                        tbl_oid, version_specific_dir, db_list[db]))
+
+        pgdata_restored_1 = self.pgdata_content(node_restored_1.data_dir)
+
+        node_restored_2 = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node_restored_2'))
+        node_restored_2.cleanup()
+        node2_tablespace = self.get_tblspace_path(node_restored_2, 'somedata')
+
+        self.restore_node(
+            backup_dir, 'node',
+            node_restored_2, options=[
+                "--db-exclude=db1",
+                "--db-exclude=db5",
+                "-T", "{0}={1}".format(
+                    node_tablespace, node2_tablespace)])
+
+        pgdata_restored_2 = self.pgdata_content(node_restored_2.data_dir)
+        self.compare_pgdata(pgdata_restored_1, pgdata_restored_2)
+
+        node_restored_2.append_conf(
+            "postgresql.auto.conf", "port = {0}".format(node_restored_2.port))
+
+        node_restored_2.slow_start()
+
+        node_restored_2.safe_psql(
+            'postgres',
+            'select 1')
+
+        try:
+            node_restored_2.safe_psql(
+                'db1',
+                'select 1')
+        except QueryException as e:
+            self.assertIn('FATAL', e.message)
+
+        try:
+            node_restored_2.safe_psql(
+                'db5',
+                'select 1')
+        except QueryException as e:
+            self.assertIn('FATAL', e.message)
+
+        with open(node_restored_2.pg_log_file, 'r') as f:
+            output = f.read()
+
+        self.assertNotIn('PANIC', output)
+
+        # Clean after yourself
+        self.del_test_dir(module_name, fname)
+
+    def test_partial_restore_include(self):
+        """
+        """
+        fname = self.id().split('.')[3]
+        backup_dir = os.path.join(self.tmp_path, module_name, fname, 'backup')
+        node = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node'),
+            initdb_params=['--data-checksums'])
+
+        self.init_pb(backup_dir)
+        self.add_instance(backup_dir, 'node', node)
+        self.set_archiving(backup_dir, 'node', node)
+        node.slow_start()
+
+        for i in range(1, 10, 1):
+            node.safe_psql(
+                'postgres',
+                'CREATE database db{0}'.format(i))
+
+        db_list_raw = node.safe_psql(
+            'postgres',
+            'SELECT to_json(a) '
+            'FROM (SELECT oid, datname FROM pg_database) a').rstrip()
+
+        db_list_splitted = db_list_raw.splitlines()
+
+        db_list = {}
+        for line in db_list_splitted:
+            line = json.loads(line)
+            db_list[line['datname']] = line['oid']
+
+        # FULL backup
+        backup_id = self.backup_node(backup_dir, 'node', node)
+        pgdata = self.pgdata_content(node.data_dir)
+
+        # restore FULL backup
+        node_restored_1 = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node_restored_1'))
+        node_restored_1.cleanup()
+
+        try:
+            self.restore_node(
+                backup_dir, 'node',
+                node_restored_1, options=[
+                    "--db-include='db1'",
+                    "--db-exclude='db2'"])
+            self.assertEqual(
+                1, 0,
+                "Expecting Error because of 'db-exclude' and 'db-include'.\n "
+                "Output: {0} \n CMD: {1}".format(
+                    self.output, self.cmd))
+        except ProbackupException as e:
+            self.assertIn(
+                "ERROR: You cannot specify '--db-include' "
+                "and '--db-exclude' together", e.message,
+                '\n Unexpected Error Message: {0}\n CMD: {1}'.format(
+                    repr(e.message), self.cmd))
+
+        self.restore_node(
+            backup_dir, 'node', node_restored_1)
+
+        pgdata_restored_1 = self.pgdata_content(node_restored_1.data_dir)
+        self.compare_pgdata(pgdata, pgdata_restored_1)
+
+        # truncate every db
+        for db in db_list:
+            # with exception below
+            if db in ['template0', 'template1', 'postgres', 'db1', 'db5']:
+                continue
+            self.truncate_every_file_in_dir(
+                os.path.join(
+                    node_restored_1.data_dir, 'base', db_list[db]))
+
+        pgdata_restored_1 = self.pgdata_content(node_restored_1.data_dir)
+
+        node_restored_2 = self.make_simple_node(
+            base_dir=os.path.join(module_name, fname, 'node_restored_2'))
+        node_restored_2.cleanup()
+
+        self.restore_node(
+            backup_dir, 'node',
+            node_restored_2, options=[
+                "--db-include=db1",
+                "--db-include=db5",
+                "--db-include=postgres"])
+
+        pgdata_restored_2 = self.pgdata_content(node_restored_2.data_dir)
+        self.compare_pgdata(pgdata_restored_1, pgdata_restored_2)
+
+        node_restored_2.append_conf(
+            "postgresql.auto.conf", "port = {0}".format(node_restored_2.port))
+        node_restored_2.slow_start()
+
+        node_restored_2.safe_psql(
+            'db1',
+            'select 1')
+
+        node_restored_2.safe_psql(
+            'db5',
+            'select 1')
+
+        node_restored_2.safe_psql(
+            'template1',
+            'select 1')
+
+        try:
+            node_restored_2.safe_psql(
+                'db2',
+                'select 1')
+        except QueryException as e:
+            self.assertIn('FATAL', e.message)
+
+        try:
+            node_restored_2.safe_psql(
+                'db10',
+                'select 1')
+        except QueryException as e:
+            self.assertIn('FATAL', e.message)
+
+        with open(node_restored_2.pg_log_file, 'r') as f:
+            output = f.read()
+
+        self.assertNotIn('PANIC', output)
+
+        # Clean after yourself
+        self.del_test_dir(module_name, fname)
+
+# partial restore
+# 0. basic test for db-exclude and db-include +
+# 1. old backup without support of partial restore
+# 2. FULL backup do not support partial restore, but incremental do
+# 3. database_map is missing for legal reasons, e.g. no permissions for pg_database
+# 4. database_map is empty for illegal reason
+# 5. database_map contain garbage