dump/restore methods

Author: zilder

testgres/testgres.py

@@ -302,6 +302,13 @@ def stop(self):
 
         return self
 
+    def restart(self):
+        """ Restarts cluster """
+        params = {"-D": self.data_dir}
+        self.pg_ctl("restart", params)
+
+        return self
+
     def reload(self):
         """Reloads config files"""
         params = {"-D": self.data_dir}
@@ -321,7 +328,7 @@ def cleanup(self):
 
         return self
 
-    def psql(self, dbname, query):
+    def psql(self, dbname, query=None, filename=None):
         """Executes a query by the psql
 
         Returns a tuple (code, stdout, stderr) in which:
@@ -331,9 +338,16 @@ def psql(self, dbname, query):
         """
         psql = self.get_bin_path("psql")
         psql_params = [
-            psql, "-XAtq", "-c", query, "-p %s" % self.port, dbname
+            psql, "-XAtq", "-p %s" % self.port, dbname
         ]
 
+        if query:
+            psql_params.extend(("-c", query))
+        elif filename:
+            psql_params.extend(("-f", filename))
+        else:
+            raise QueryException('Query or filename must be provided')
+
         # start psql process
         process = subprocess.Popen(
             psql_params,
@@ -356,6 +370,35 @@ def safe_psql(self, dbname, query):
             raise ClusterException("psql failed:\n" + err)
         return out
 
+    def dump(self, dbname, filename):
+        """Invoke pg_dump and exports database to a file as an sql script"""
+        path = self.base_dir + "/" + filename
+        params = [
+            self.get_bin_path("pg_dump"),
+            "-p %s" % self.port,
+            "-f", path,
+            dbname
+        ]
+
+        with open(self.error_filename, "a") as file_err:
+            ret = subprocess.call(params, stderr=file_err)
+            if ret:
+                raise ClusterException("Dump creation failed")
+
+    def restore(self, dbname, filename, node=None):
+        """ Restore database from dump file
+
+        dbname     name of database we're restoring data to
+        filename   path to the dump file
+        node       the cluster node where the dump was taken from (if node
+                   wasn't specified then assume the same cluster)
+        """
+        if not node:
+            node = self
+
+        path = node.base_dir + "/" + filename
+        self.psql(dbname, filename=path)
+
     def poll_query_until(self, dbname, query):
         """Runs a query once a second until it returs True"""
         max_attemps = 60

testgres/tests/test_simple.py

@@ -18,6 +18,7 @@ def test_start_stop(self):
 		self.assertEqual(res[0][0], 1)
 		node.stop()
 
+	@unittest.skip("demo")
 	def test_backup_and_replication(self):
 		node = get_new_node('test')
 		replica = get_new_node('repl')
@@ -51,5 +52,22 @@ def test_backup_and_replication(self):
 		node.stop()
 		replica.stop()
 
+	def test_dump(self):
+		node = get_new_node('test')
+		node.init().start()
+		node.safe_psql(
+			'postgres',
+			'create table abc as '
+			'select g as a, g as b from generate_series(1, 10) as g'
+		)
+		node.psql('postgres', 'create database test')
+		node.dump('postgres', 'test.sql')
+		node.restore('test', 'test.sql')
+		self.assertEqual(
+			node.psql('postgres', 'select * from abc'),
+			node.psql('test', 'select * from abc'),
+		)
+		node.stop()
+
 if __name__ == '__main__':
 	unittest.main()