From e40943c25f48eeb06971ed404002e4e284b0065b Mon Sep 17 00:00:00 2001
From: Rushabh Lathia <rushabh.lathia@enterprisedb.com>
Date: Wed, 5 Feb 2025 15:38:21 +0530
Subject: [PATCH 3/3] Support pg_dump to dump NOT VALID named NOT NULL
 constraints.

Dump the INVALID NOT NULL constraint as separate from table schema,
just like CHECK constraints.
---
 src/bin/pg_dump/pg_dump.c | 168 ++++++++++++++++++++++++++++++++++++++++++----
 src/bin/pg_dump/pg_dump.h |   1 +
 2 files changed, 156 insertions(+), 13 deletions(-)

diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index 02e1fdf..319eeab 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -8745,6 +8745,7 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 	PQExpBuffer q = createPQExpBuffer();
 	PQExpBuffer tbloids = createPQExpBuffer();
 	PQExpBuffer checkoids = createPQExpBuffer();
+	PQExpBuffer invalidnotnulloids = createPQExpBuffer();
 	PGresult   *res;
 	int			ntups;
 	int			curtblindx;
@@ -8761,6 +8762,7 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 	int			i_attlen;
 	int			i_attalign;
 	int			i_attislocal;
+	int			i_notnull_valid;
 	int			i_notnull_name;
 	int			i_notnull_noinherit;
 	int			i_notnull_islocal;
@@ -8857,11 +8859,13 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 	 */
 	if (fout->remoteVersion >= 180000)
 		appendPQExpBufferStr(q,
+							 "CASE WHEN a.attnotnull = 'i' THEN 'f' ELSE 't' END AS notnull_valid,\n"
 							 "co.conname AS notnull_name,\n"
 							 "co.connoinherit AS notnull_noinherit,\n"
 							 "co.conislocal AS notnull_islocal,\n");
 	else
 		appendPQExpBufferStr(q,
+							 "'t' AS notnull_valid,\n"
 							 "CASE WHEN a.attnotnull THEN '' ELSE NULL END AS notnull_name,\n"
 							 "false AS notnull_noinherit,\n"
 							 "a.attislocal AS notnull_islocal,\n");
@@ -8936,6 +8940,7 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 	i_attlen = PQfnumber(res, "attlen");
 	i_attalign = PQfnumber(res, "attalign");
 	i_attislocal = PQfnumber(res, "attislocal");
+	i_notnull_valid = PQfnumber(res, "notnull_valid");
 	i_notnull_name = PQfnumber(res, "notnull_name");
 	i_notnull_noinherit = PQfnumber(res, "notnull_noinherit");
 	i_notnull_islocal = PQfnumber(res, "notnull_islocal");
@@ -8955,6 +8960,7 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 	 * r is handled by the inner loop.
 	 */
 	curtblindx = -1;
+	appendPQExpBufferChar(invalidnotnulloids, '{');
 	for (int r = 0; r < ntups;)
 	{
 		Oid			attrelid = atooid(PQgetvalue(res, r, i_attrelid));
@@ -9003,6 +9009,7 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 		tbinfo->attcompression = (char *) pg_malloc(numatts * sizeof(char));
 		tbinfo->attfdwoptions = (char **) pg_malloc(numatts * sizeof(char *));
 		tbinfo->attmissingval = (char **) pg_malloc(numatts * sizeof(char *));
+		tbinfo->notnull_valid = (bool *) pg_malloc(numatts * sizeof(bool));
 		tbinfo->notnull_constrs = (char **) pg_malloc(numatts * sizeof(char *));
 		tbinfo->notnull_noinh = (bool *) pg_malloc(numatts * sizeof(bool));
 		tbinfo->notnull_islocal = (bool *) pg_malloc(numatts * sizeof(bool));
@@ -9029,12 +9036,28 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 			tbinfo->attlen[j] = atoi(PQgetvalue(res, r, i_attlen));
 			tbinfo->attalign[j] = *(PQgetvalue(res, r, i_attalign));
 			tbinfo->attislocal[j] = (PQgetvalue(res, r, i_attislocal)[0] == 't');
+			tbinfo->notnull_valid[j] = (PQgetvalue(res, r, i_notnull_valid)[0] == 't');
 
-			/* Handle not-null constraint name and flags */
-			determineNotNullFlags(fout, res, r,
-								  tbinfo, j,
-								  i_notnull_name, i_notnull_noinherit,
-								  i_notnull_islocal);
+			/*
+			 * Dump the invalid NOT NULL constrint like the Check constraints
+			 */
+			if (tbinfo->notnull_valid[j])
+				/* Handle not-null constraint name and flags */
+				determineNotNullFlags(fout, res, r,
+									  tbinfo, j,
+									  i_notnull_name, i_notnull_noinherit,
+									  i_notnull_islocal);
+			else if(!PQgetisnull(res, r, i_notnull_name))
+			{
+				/*
+				 * Add the entry into invalidnotnull list so NOT NULL constraint get
+				 * dump as separate constrints.
+				 */
+				if (invalidnotnulloids->len > 1) /* do we have more than the '{'? */
+					appendPQExpBufferChar(invalidnotnulloids, ',');
+				appendPQExpBuffer(invalidnotnulloids, "%u", tbinfo->dobj.catId.oid);
+
+			}
 
 			tbinfo->attoptions[j] = pg_strdup(PQgetvalue(res, r, i_attoptions));
 			tbinfo->attcollation[j] = atooid(PQgetvalue(res, r, i_attcollation));
@@ -9056,6 +9079,7 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 	}
 
 	PQclear(res);
+	appendPQExpBufferChar(invalidnotnulloids, '}');
 
 	/*
 	 * Now get info about column defaults.  This is skipped for a data-only
@@ -9319,6 +9343,128 @@ getTableAttrs(Archive *fout, TableInfo *tblinfo, int numTables)
 		PQclear(res);
 	}
 
+	/*
+	 * Get info about table INVALID NOT NULL constraints.  This is skipped for a
+	 * data-only dump, as it is only needed for table schemas.
+	 */
+	if (dopt->dumpSchema && invalidnotnulloids->len > 2)
+	{
+		ConstraintInfo *constrs;
+		int			numConstrs;
+		int			i_tableoid;
+		int			i_oid;
+		int			i_conrelid;
+		int			i_conname;
+		int			i_consrc;
+		int			i_conislocal;
+		int			i_convalidated;
+
+		pg_log_info("finding table check constraints");
+
+		resetPQExpBuffer(q);
+		appendPQExpBuffer(q,
+						  "SELECT c.tableoid, c.oid, conrelid, conname, "
+						  "pg_catalog.pg_get_constraintdef(c.oid) AS consrc, "
+						  "conislocal, convalidated "
+						  "FROM unnest('%s'::pg_catalog.oid[]) AS src(tbloid)\n"
+						  "JOIN pg_catalog.pg_constraint c ON (src.tbloid = c.conrelid)\n"
+						  "WHERE contype = 'n' AND convalidated = 'f'"
+						  "ORDER BY c.conrelid, c.conname",
+						  invalidnotnulloids->data);
+
+		res = ExecuteSqlQuery(fout, q->data, PGRES_TUPLES_OK);
+
+		numConstrs = PQntuples(res);
+		constrs = (ConstraintInfo *) pg_malloc(numConstrs * sizeof(ConstraintInfo));
+
+		i_tableoid = PQfnumber(res, "tableoid");
+		i_oid = PQfnumber(res, "oid");
+		i_conrelid = PQfnumber(res, "conrelid");
+		i_conname = PQfnumber(res, "conname");
+		i_consrc = PQfnumber(res, "consrc");
+		i_conislocal = PQfnumber(res, "conislocal");
+		i_convalidated = PQfnumber(res, "convalidated");
+
+		/* As above, this loop iterates once per table, not once per row */
+		curtblindx = -1;
+		for (int j = 0; j < numConstrs;)
+		{
+			Oid			conrelid = atooid(PQgetvalue(res, j, i_conrelid));
+			TableInfo  *tbinfo = NULL;
+			int			numcons;
+
+			/* Count rows for this table */
+			for (numcons = 1; numcons < numConstrs - j; numcons++)
+				if (atooid(PQgetvalue(res, j + numcons, i_conrelid)) != conrelid)
+					break;
+
+			/*
+			 * Locate the associated TableInfo; we rely on tblinfo[] being in
+			 * OID order.
+			 */
+			while (++curtblindx < numTables)
+			{
+				tbinfo = &tblinfo[curtblindx];
+				if (tbinfo->dobj.catId.oid == conrelid)
+					break;
+			}
+			if (curtblindx >= numTables)
+				pg_fatal("unrecognized table OID %u", conrelid);
+
+
+			tbinfo->checkexprs = constrs + j;
+
+			for (int c = 0; c < numcons; c++, j++)
+			{
+				bool		validated = PQgetvalue(res, j, i_convalidated)[0] == 't';
+
+				constrs[j].dobj.objType = DO_CONSTRAINT;
+				constrs[j].dobj.catId.tableoid = atooid(PQgetvalue(res, j, i_tableoid));
+				constrs[j].dobj.catId.oid = atooid(PQgetvalue(res, j, i_oid));
+				AssignDumpId(&constrs[j].dobj);
+				constrs[j].dobj.name = pg_strdup(PQgetvalue(res, j, i_conname));
+				constrs[j].dobj.namespace = tbinfo->dobj.namespace;
+				constrs[j].contable = tbinfo;
+				constrs[j].condomain = NULL;
+				constrs[j].contype = 'n';
+				constrs[j].condef = pg_strdup(PQgetvalue(res, j, i_consrc));
+				constrs[j].confrelid = InvalidOid;
+				constrs[j].conindex = 0;
+				constrs[j].condeferrable = false;
+				constrs[j].condeferred = false;
+				constrs[j].conislocal = (PQgetvalue(res, j, i_conislocal)[0] == 't');
+
+				/*
+				 * An unvalidated constraint needs to be dumped separately, so
+				 * that potentially-violating existing data is loaded before
+				 * the constraint.
+				 */
+				constrs[j].separate = !validated;
+
+				constrs[j].dobj.dump = tbinfo->dobj.dump;
+
+				/*
+				 * Mark the constraint as needing to appear before the table
+				 * --- this is so that any other dependencies of the
+				 * constraint will be emitted before we try to create the
+				 * table.  If the constraint is to be dumped separately, it
+				 * will be dumped after data is loaded anyway, so don't do it.
+				 * (There's an automatic dependency in the opposite direction
+				 * anyway, so don't need to add one manually here.)
+				 */
+				if (!constrs[j].separate)
+					addObjectDependency(&tbinfo->dobj,
+										constrs[j].dobj.dumpId);
+
+				/*
+				 * We will detect later whether the constraint must be split
+				 * out from the table definition.
+				 */
+			}
+		}
+		PQclear(res);
+	}
+
 	destroyPQExpBuffer(q);
 	destroyPQExpBuffer(tbloids);
 	destroyPQExpBuffer(checkoids);
@@ -16136,7 +16282,8 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 					 * defined, or if binary upgrade.  (In the latter case, we
 					 * reset conislocal below.)
 					 */
-					print_notnull = (tbinfo->notnull_constrs[j] != NULL &&
+					print_notnull = (tbinfo->notnull_valid[j] &&
+									 tbinfo->notnull_constrs[j] != NULL &&
 									 (tbinfo->notnull_islocal[j] ||
 									  dopt->binary_upgrade ||
 									  tbinfo->ispartition));
@@ -16195,11 +16342,6 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 											  tbinfo->attrdefs[j]->adef_expr);
 					}
 
-					print_notnull = (tbinfo->notnull_constrs[j] != NULL &&
-									 (tbinfo->notnull_islocal[j] ||
-									  dopt->binary_upgrade ||
-									  tbinfo->ispartition));
-
 					if (print_notnull)
 					{
 						if (tbinfo->notnull_constrs[j][0] == '\0')
@@ -17521,9 +17663,9 @@ dumpConstraint(Archive *fout, const ConstraintInfo *coninfo)
 									  .createStmt = q->data,
 									  .dropStmt = delq->data));
 	}
-	else if (coninfo->contype == 'c' && tbinfo)
+	else if ((coninfo->contype == 'c' || coninfo->contype == 'n') && tbinfo)
 	{
-		/* CHECK constraint on a table */
+		/* CHECK/INVALID_NOTNULL constraint on a table */
 
 		/* Ignore if not to be dumped separately, or if it was inherited */
 		if (coninfo->separate && coninfo->conislocal)
diff --git a/src/bin/pg_dump/pg_dump.h b/src/bin/pg_dump/pg_dump.h
index 7139c88..c4e05bd 100644
--- a/src/bin/pg_dump/pg_dump.h
+++ b/src/bin/pg_dump/pg_dump.h
@@ -356,6 +356,7 @@ typedef struct _tableInfo
 	char	   *attcompression; /* per-attribute compression method */
 	char	  **attfdwoptions;	/* per-attribute fdw options */
 	char	  **attmissingval;	/* per attribute missing value */
+	bool       *notnull_valid;  /* NOT NULL status */
 	char	  **notnull_constrs;	/* NOT NULL constraint names. If null,
 									 * there isn't one on this column. If
 									 * empty string, unnamed constraint
-- 
1.8.3.1

