replication

Author: zilder

testgres/__init__.py

@@ -1 +1 @@
-from testgres import PostgresNode, get_new_node, clean_all
+from testgres import PostgresNode, get_new_node, clean_all, stop_all

testgres/testgres.py

@@ -28,6 +28,7 @@
 import pwd
 import tempfile
 import shutil
+import time
 
 # Try to use psycopg2 by default. If psycopg2 isn"t available then use 
 # pg8000 which is slower but much more portable because uses only
@@ -45,13 +46,14 @@
 last_assigned_port = int(random.random() * 16384) + 49152;
 
 
-class ClusterException(Exception):
-    pass
+class ClusterException(Exception): pass
+class QueryException(Exception): pass
 
 
 class PostgresNode:
     def __init__(self, name, port):
         self.name = name
+        self.host = '127.0.0.1'
         self.port = port
         self.base_dir = tempfile.mkdtemp()
         os.makedirs(self.logs_dir)
@@ -75,6 +77,11 @@ def output_filename(self):
     def error_filename(self):
         return "%s/stderr.log" % self.logs_dir
 
+    @property
+    def connstr(self):
+        return "port=%s" % self.port    
+        # return "port=%s host=%s" % (self.port, self.host)
+
     def load_pg_config(self):
         """ Loads pg_config output into dict """
         pg_config = os.environ.get("PG_CONFIG") \
@@ -116,6 +123,10 @@ def init(self, allows_streaming=False):
                 "fsync = off\n"
                 "log_statement = all\n"
                 "port = %s\n" % self.port)
+            conf.write(
+                # "unix_socket_directories = '%s'\n"
+                # "listen_addresses = ''\n";)
+                "listen_addresses = '%s'\n" % self.host)
 
             if allows_streaming:
                 conf.write(
@@ -129,7 +140,7 @@ def init(self, allows_streaming=False):
                     "max_connections = 10\n")
         self.set_replication_conf()
 
-    def init_from_backup(self, root_node, backup_name):
+    def init_from_backup(self, root_node, backup_name, has_streaming=False, hba_permit_replication=True):
         """Initializes cluster from backup, made by another node"""
 
         # Copy data from backup
@@ -142,11 +153,25 @@ def init_from_backup(self, root_node, backup_name):
             "postgresql.conf",
             "port = %s" % self.port
         )
+        # Enable streaming
+        if hba_permit_replication:
+            self.set_replication_conf()
+        if has_streaming:
+            self.enable_streaming(root_node)
 
     def set_replication_conf(self):
         hba_conf = "%s/pg_hba.conf" % self.data_dir
         with open(hba_conf, "a") as conf:
             conf.write("local replication all trust\n")
+            # conf.write("host replication all 127.0.0.1/32 trust\n")
+
+    def enable_streaming(self, root_node):
+        config_name = "%s/recovery.conf" % self.data_dir
+        with open(config_name, "a") as conf:
+            conf.write(
+                "primary_conninfo='%s application_name=%s'\n"
+                "standby_mode=on\n"
+                % (root_node.connstr, self.name))
 
     def append_conf(self, filename, string):
         """Appends line to a config file like "postgresql.conf"
@@ -254,6 +279,22 @@ def safe_psql(self, dbname, query):
             raise ClusterException("psql failed:\n" + err)
         return out
 
+    def poll_query_until(self, dbname, query):
+        """Runs a query once a second until it returs True"""
+        max_attemps = 60
+        attemps = 0
+
+        while attemps < max_attemps:
+            ret = self.safe_psql(dbname, query)
+
+            # TODO: fix psql so that it returns result without newline
+            if ret == "t\n":
+                return
+
+            time.sleep(1)
+            attemps += 1
+        raise QueryException("Timeout while waiting for query to return True")
+
     def execute(self, dbname, query):
         """Executes the query and returns all rows"""
         connection = pglib.connect(
@@ -327,3 +368,8 @@ def clean_all():
     for node in registered_nodes:
         node.cleanup()
     registered_nodes = []
+
+def stop_all():
+    global registered_nodes
+    for node in registered_nodes:
+        node.stop()

testgres/tests/test_simple.py

@@ -1,11 +1,12 @@
 import unittest
-from testgres import get_new_node, clean_all
+import time
+from testgres import get_new_node, clean_all, stop_all
 
 class SimpleTest(unittest.TestCase):
 
 	def teardown(self):
 		# clean_all()
-		pass
+		stop_all()
 
 	@unittest.skip("demo")
 	def test_start_stop(self):
@@ -17,7 +18,7 @@ def test_start_stop(self):
 		self.assertEqual(res[0][0], 1)
 		node.stop()
 
-	def test_backup(self):
+	def test_backup_and_replication(self):
 		node = get_new_node('test')
 		replica = get_new_node('repl')
 
@@ -27,13 +28,28 @@ def test_backup(self):
 		node.psql('postgres', 'insert into abc values (1, 2)')
 		node.backup('my_backup')
 
-		replica.init_from_backup(node, 'my_backup')
+		replica.init_from_backup(node, 'my_backup', has_streaming=True)
 		replica.start()
 		res = replica.execute('postgres', 'select * from abc')
 		self.assertEqual(len(res), 1)
 		self.assertEqual(res[0], (1, 2))
 
+		# Insert into master node
+		node.psql('postgres', 'insert into abc values (3, 4)')
+		# Wait until data syncronizes
+		node.poll_query_until(
+			'postgres',
+			'SELECT pg_current_xlog_location() <= write_location '
+			'FROM pg_stat_replication WHERE application_name = \'%s\''
+			% replica.name)
+		# time.sleep(0.5)
+		# Check that this record was exported to replica
+		res = replica.execute('postgres', 'select * from abc')
+		self.assertEqual(len(res), 2)
+		self.assertEqual(res[1], (3, 4))
+
 		node.stop()
+		replica.stop()
 
 if __name__ == '__main__':
 	unittest.main()