Fix gp_update_ao_master_stats for gpdbrestore with full backup.

When doing a restore of a specific table or a specific schema (-T or -S)
with a full backup, it will update all statistics for all AO/CO tables; it should
update only the specified table/schema.

Moved gp_update_ao_master_stats from the C side to the python code.
gpdbrestore -T and -S will now update AO/CO stats through the python code

Authors: Marbin Tan & Chris Hajas

gpMgmt/bin/gppylib/operations/restore.py

@@ -45,13 +45,30 @@ POST_DATA_SUFFIX = '_post_data'
 # TODO: use CLI-agnostic custom exceptions instead of ExceptionNoStackTraceNeeded
 
 def update_ao_stat_func(conn, ao_schema, ao_table, counter, batch_size):
-    qry = "SELECT * FROM gp_update_ao_master_stats('%s.%s')" % (pg.escape_string(escapeDoubleQuoteInSQLString(ao_schema)), 
+    qry = "SELECT * FROM gp_update_ao_master_stats('%s.%s')" % (pg.escape_string(escapeDoubleQuoteInSQLString(ao_schema)),
                                                                 pg.escape_string(escapeDoubleQuoteInSQLString(ao_table)))
     rows = execSQLForSingleton(conn, qry)
     if counter % batch_size == 0:
         conn.commit()
 
-def update_ao_statistics(master_port, dbname, restored_tables):
+def generate_restored_tables(results, restored_tables, restored_schema, restore_all):
+    restored_ao_tables = set()
+
+    for (tbl, sch) in results:
+        if restore_all:
+            restored_ao_tables.add((sch, tbl))
+        elif sch in restored_schema:
+            restored_ao_tables.add((sch, tbl))
+        else:
+            tblname = '%s.%s' % (sch, tbl)
+            if tblname in restored_tables:
+                restored_ao_tables.add((sch, tbl))
+
+    return restored_ao_tables
+
+def update_ao_statistics(master_port, dbname, restored_tables, restored_schema=[], restore_all=False):
+    # Restored schema is different from restored tables as restored schema
+    # updates all tables within that schema.
     qry = """SELECT c.relname,n.nspname
              FROM pg_class c, pg_namespace n
              WHERE c.relnamespace=n.oid
@@ -60,12 +77,13 @@ def update_ao_statistics(master_port, dbname, restored_tables):
     conn = None
     counter = 1
     restored_ao_tables = set()
+
     try:
         results = execute_sql(qry, master_port, dbname)
-        for (tbl, sch) in results:
-            tblname = '%s.%s' % (sch, tbl)
-            if tblname in restored_tables:
-                restored_ao_tables.add((sch, tbl))
+        restored_ao_tables = generate_restored_tables(results,
+                                                      restored_tables,
+                                                      restored_schema,
+                                                      restore_all)
 
         if len(restored_ao_tables) == 0:
             logger.info("No AO/CO tables restored, skipping statistics update...")
@@ -588,6 +606,17 @@ class RestoreDatabase(Operation):
                 cmd.run(validateAfter=False)
                 self._process_result(cmd)
 
+            restore_all=False
+            if not self.no_ao_stats:
+                logger.info("Updating AO/CO statistics on master")
+                # If we don't have a filter for table and schema, then we must
+                # be doing a full restore.
+                if len(self.schema_level_restore_list) == 0 and len(self.restore_tables) == 0:
+                    restore_all=True
+                update_ao_statistics(self.master_port, restore_db, self.restore_tables,
+                                     restored_schema=self.schema_level_restore_list, restore_all=restore_all,
+                                    )
+
         if not self.metadata_only:
             if (not self.no_analyze) and (len(self.restore_tables) == 0):
                 self._analyze(restore_db, self.master_port)
@@ -848,7 +877,7 @@ class RestoreDatabase(Operation):
 
             logger.info("Dropping Database %s" % restore_db)
             if count == 1:
-                cmd = Command(name='drop database %s' % restore_db, 
+                cmd = Command(name='drop database %s' % restore_db,
                               cmdStr='dropdb %s -p %s' % (checkAndAddEnclosingDoubleQuote(shellEscape(restore_db)), master_port))
                 cmd.run(validateAfter=True)
             logger.info("Dropped Database %s" % restore_db)
@@ -978,7 +1007,7 @@ class RestoreDatabase(Operation):
         return restore_line
 
     def _build_post_data_schema_only_restore_line(self, restore_timestamp, restore_db, compress, master_port,
-                                                  table_filter_file, full_restore_with_filter, 
+                                                  table_filter_file, full_restore_with_filter,
                                                   change_schema_file=None, schema_level_restore_file=None):
         user = getpass.getuser()
         hostname = socket.gethostname()    # TODO: can this just be localhost? bash was using `hostname`
@@ -1023,7 +1052,7 @@ class RestoreDatabase(Operation):
         return restore_line
 
     def _build_schema_only_restore_line(self, restore_timestamp, restore_db, compress, master_port,
-                                        metadata_filename, table_filter_file, full_restore_with_filter, 
+                                        metadata_filename, table_filter_file, full_restore_with_filter,
                                         change_schema_file=None, schema_level_restore_file=None):
         user = getpass.getuser()
         hostname = socket.gethostname()    # TODO: can this just be localhost? bash was using `hostname`
@@ -1212,7 +1241,7 @@ def validate_tablenames(table_list, master_data_dir, backup_dir, dump_dir, dump_
         if table not in dumped_tables:
             raise Exception("Table %s not found in backup" % table)
 
-    return restore_table_list, schema_level_restore_list 
+    return restore_table_list, schema_level_restore_list
 
 class ValidateRestoreTables(Operation):
     def __init__(self, restore_tables, restore_db, master_port):

src/bin/pg_dump/cdb/cdb_restore.c

@@ -207,11 +207,6 @@ main(int argc, char **argv)
 			mpp_err_msg(logInfo, progname, "All remote %s programs are finished.\n", pszAgent);
 		}
 
-		/*
-		 * If any AO table data was restored, update the master AO statistics
-		 */
-		if (bAoStats)
-			updateAppendOnlyStats(pConn);
 	}
 
 	/*
@@ -1393,112 +1388,6 @@ restoreMaster(InputOptions * pInputOpts,
 	return rtn;
 }
 
-/*
- * updateAppendOnlyStats
- *
- * for every append only that exists in the database (whether we just restored
- * it or not) update its statistics on the master - note that this must not be
- * made with utility mode connection, but rather a regular connection.
- */
-void
-updateAppendOnlyStats(PGconn *pConn)
-{
-	/* query that gets all ao tables in the database */
-	PQExpBuffer get_query = createPQExpBuffer();
-	PGresult   *get_res;
-	int			get_ntups = 0;
-
-	/* query that updates all ao tables in the database */
-	PQExpBuffer update_query = createPQExpBuffer();
-	PGresult   *update_res = NULL;
-	int			update_ntups = 0;
-
-	/* misc */
-	int			i_tablename = 0;
-	int			i_schemaname = 0;
-	int			i = 0;
-
-	mpp_err_msg(logInfo, progname, "updating Append Only table statistics\n");
-
-	/* Fetch all Append Only tables */
-	appendPQExpBuffer(get_query,
-					  "SELECT c.relname,n.nspname "
-					  "FROM pg_class c, pg_namespace n "
-					  "WHERE c.relnamespace=n.oid "
-					  "AND (c.relstorage='a' OR c.relstorage='c')");
-
-	get_res = PQexec(pConn, get_query->data);
-
-	if (!get_res || PQresultStatus(get_res) != PGRES_TUPLES_OK)
-	{
-		const char *err;
-
-		if (get_res)
-			err = PQresultErrorMessage(get_res);
-		else
-			err = PQerrorMessage(pConn);
-
-		mpp_err_msg(logWarn, progname,
-					"Failed to get Append Only tables for updating stats. "
-					"error was: %s\n", err);
-
-	}
-	else
-	{
-		/* A-OK */
-		get_ntups = PQntuples(get_res);
-		i_tablename = PQfnumber(get_res, "relname");
-		i_schemaname = PQfnumber(get_res, "nspname");
-
-		for (i = 0; i < get_ntups; i++)
-		{
-			char	   *tablename = PQgetvalue(get_res, i, i_tablename);
-			char	   *schemaname = PQgetvalue(get_res, i, i_schemaname);
-
-			resetPQExpBuffer(update_query);
-			appendPQExpBuffer(update_query,
-							  "SELECT * "
-							  "FROM gp_update_ao_master_stats('%s.%s')",
-							  schemaname, tablename);
-
-			update_res = PQexec(pConn, update_query->data);
-
-			if (!update_res || PQresultStatus(update_res) != PGRES_TUPLES_OK)
-			{
-				const char *err;
-
-				if (update_res)
-					err = PQresultErrorMessage(update_res);
-				else
-					err = PQerrorMessage(pConn);
-
-				mpp_err_msg(logWarn, progname,
-							"Failed to update Append Only table stats. error "
-						"was: %s, query was: %s\n", err, update_query->data);
-			}
-			else
-			{
-				/* A-OK */
-				update_ntups = PQntuples(update_res);
-				if (update_ntups != 1)
-					mpp_err_msg(logWarn, progname,
-								"Expected 1 tuple, got %d. query was: %s\n",
-								update_ntups, update_query->data);
-			}
-		}
-	}
-
-	/* clean up */
-	if (update_res)
-		PQclear(update_res);
-
-	if (get_res)
-		PQclear(get_res);
-
-	destroyPQExpBuffer(get_query);
-	destroyPQExpBuffer(update_query);
-
-}
 
 /*
  * spinOffThreads: This function deals with all the threads that drive the backend restores.