gprecoverseg: add tablespace test

Verify tablespaces work for both a incremental and full recovery.
Co-authored-by: NJacob Champion <pchampion@pivotal.io>
Authored-by: NKalen Krempely <kkrempely@pivotal.io>

gpMgmt/test/behave/mgmt_utils/environment.py

@@ -76,6 +76,10 @@ def before_scenario(context, scenario):
 
 
 def after_scenario(context, scenario):
+    if 'tablespaces' in context:
+        for tablespace in context.tablespaces.values():
+            tablespace.cleanup()
+
     tags_to_skip = ['gpexpand', 'gpaddmirrors', 'gpstate', 'gpinitstandby']
     if set(context.feature.tags).intersection(tags_to_skip):
         return

gpMgmt/test/behave/mgmt_utils/gprecoverseg.feature

 @gprecoverseg
 Feature: gprecoverseg tests
 
+    Scenario: incremental recovery works with tablespaces
+        Given the database is running
+          And a tablespace is created with data
+          And user kills a primary postmaster process
+          And user can start transactions
+         When the user runs "gprecoverseg -a"
+         Then gprecoverseg should return a return code of 0
+          And the segments are synchronized
+          And the tablespace is valid
+
+        Given another tablespace is created with data
+         When the user runs "gprecoverseg -ra"
+         Then gprecoverseg should return a return code of 0
+          And the segments are synchronized
+          And the tablespace is valid
+          And the other tablespace is valid
+
+    Scenario: full recovery works with tablespaces
+        Given the database is running
+          And a tablespace is created with data
+          And user kills a primary postmaster process
+          And user can start transactions
+         When the user runs "gprecoverseg -a -F"
+         Then gprecoverseg should return a return code of 0
+          And the segments are synchronized
+          And the tablespace is valid
+
+        Given another tablespace is created with data
+         When the user runs "gprecoverseg -ra"
+         Then gprecoverseg should return a return code of 0
+          And the segments are synchronized
+          And the tablespace is valid
+          And the other tablespace is valid
+
     Scenario: gprecoverseg should not output bootstrap error on success
         Given the database is running
         And user kills a primary postmaster process

gpMgmt/test/behave/mgmt_utils/steps/tablespace_mgmt_utils.py

+import shutil
+import tempfile
+
+from behave import given, when, then
+from pygresql import pg
+
+from gppylib.db import dbconn
+
+
+class Tablespace:
+    def __init__(self, name):
+        self.name = name
+        self.path = tempfile.mkdtemp()
+        self.dbname = 'tablespace_db_%s' % name
+        self.table_counter = 0
+        self.initial_data = None
+
+        with dbconn.connect(dbconn.DbURL()) as conn:
+            db = pg.DB(conn)
+            db.query("CREATE TABLESPACE %s LOCATION '%s'" % (self.name, self.path))
+            db.query("CREATE DATABASE %s TABLESPACE %s" % (self.dbname, self.name))
+
+        with dbconn.connect(dbconn.DbURL(dbname=self.dbname)) as conn:
+            db = pg.DB(conn)
+            db.query("CREATE TABLE tbl (i int) DISTRIBUTED RANDOMLY")
+            db.query("INSERT INTO tbl VALUES (GENERATE_SERIES(0, 25))")
+            # save the distributed data for later verification
+            self.initial_data = db.query("SELECT gp_segment_id, i FROM tbl").getresult()
+
+    def cleanup(self):
+        with dbconn.connect(dbconn.DbURL()) as conn:
+            db = pg.DB(conn)
+            db.query("DROP DATABASE IF EXISTS %s" % self.dbname)
+            db.query("DROP TABLESPACE IF EXISTS %s" % self.name)
+
+        shutil.rmtree(self.path)
+
+    def verify(self):
+        """
+        Verify tablespace functionality by ensuring the tablespace can be
+        written to, read from, and the initial data is still correctly
+        distributed.
+        """
+        with dbconn.connect(dbconn.DbURL(dbname=self.dbname)) as conn:
+            db = pg.DB(conn)
+            data = db.query("SELECT gp_segment_id, i FROM tbl").getresult()
+
+            # verify that we can still write to the tablespace
+            self.table_counter += 1
+            db.query("CREATE TABLE tbl_%s (i int) DISTRIBUTED RANDOMLY" % self.table_counter)
+            db.query("INSERT INTO tbl_%s VALUES (GENERATE_SERIES(0, 25))" % self.table_counter)
+
+        if sorted(data) != sorted(self.initial_data):
+            raise Exception("Tablespace data is not identically distributed. Expected:\n%r\n but found:\n%r" % (
+                sorted(self.initial_data), sorted(data)))
+
+
+@given('a tablespace is created with data')
+def impl(context):
+    _create_tablespace_with_data(context, "outerspace")
+
+
+@given('another tablespace is created with data')
+def impl(context):
+    _create_tablespace_with_data(context, "myspace")
+
+
+def _create_tablespace_with_data(context, name):
+    if 'tablespaces' not in context:
+        context.tablespaces = {}
+    context.tablespaces[name] = Tablespace(name)
+
+
+@then('the tablespace is valid')
+def impl(context):
+    context.tablespaces["outerspace"].verify()
+
+
+@then('the other tablespace is valid')
+def impl(context):
+    context.tablespaces["myspace"].verify()