sangli00
diff --git a/‎contrib/dblink/dblink.c
Lines changed: 52 additions & 0 deletions b/‎contrib/dblink/dblink.c
Lines changed: 52 additions & 0 deletions
diff --git a/‎contrib/dblink/expected/dblink.out
Lines changed: 9 additions & 0 deletions b/‎contrib/dblink/expected/dblink.out
Lines changed: 9 additions & 0 deletions
diff --git a/‎contrib/dblink/sql/dblink.sql
Lines changed: 6 additions & 0 deletions b/‎contrib/dblink/sql/dblink.sql
Lines changed: 6 additions & 0 deletions
@@ -76,6 +76,7 @@ static void remove_res_ptr(dblink_results * results);
 static char *generate_relation_name(Oid relid);
 static char *connstr_strip_password(const char *connstr);
 static void dblink_security_check(PGconn *conn, const char *connstr);
+static int get_nondropped_natts(Oid relid);
 
 /* Global */
 List	   *res_id = NIL;
@@ -1100,6 +1101,7 @@ dblink_build_sql_insert(PG_FUNCTION_ARGS)
 	int16		typlen;
 	bool		typbyval;
 	char		typalign;
+	int			nondropped_natts;
 
 	relname_text = PG_GETARG_TEXT_P(0);
 
@@ -1123,6 +1125,14 @@ dblink_build_sql_insert(PG_FUNCTION_ARGS)
 	if (pknumatts == 0)
 		elog(ERROR, "dblink_build_sql_insert: number of key attributes must be > 0.");
 
+	/*
+	 * ensure we don't ask for more pk attributes than we have
+	 * non-dropped columns
+	 */
+	nondropped_natts = get_nondropped_natts(relid);
+	if (pknumatts > nondropped_natts)
+		elog(ERROR, "number of primary key fields exceeds number of specified relation attributes");
+
 	src_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(3);
 	tgt_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(4);
 
@@ -1239,6 +1249,7 @@ dblink_build_sql_delete(PG_FUNCTION_ARGS)
 	int16		typlen;
 	bool		typbyval;
 	char		typalign;
+	int			nondropped_natts;
 
 	relname_text = PG_GETARG_TEXT_P(0);
 
@@ -1262,6 +1273,14 @@ dblink_build_sql_delete(PG_FUNCTION_ARGS)
 	if (pknumatts == 0)
 		elog(ERROR, "dblink_build_sql_insert: number of key attributes must be > 0.");
 
+	/*
+	 * ensure we don't ask for more pk attributes than we have
+	 * non-dropped columns
+	 */
+	nondropped_natts = get_nondropped_natts(relid);
+	if (pknumatts > nondropped_natts)
+		elog(ERROR, "number of primary key fields exceeds number of specified relation attributes");
+
 	tgt_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(3);
 
 	/*
@@ -1356,6 +1375,7 @@ dblink_build_sql_update(PG_FUNCTION_ARGS)
 	int16		typlen;
 	bool		typbyval;
 	char		typalign;
+	int			nondropped_natts;
 
 	relname_text = PG_GETARG_TEXT_P(0);
 
@@ -1379,6 +1399,14 @@ dblink_build_sql_update(PG_FUNCTION_ARGS)
 	if (pknumatts == 0)
 		elog(ERROR, "dblink_build_sql_insert: number of key attributes must be > 0.");
 
+	/*
+	 * ensure we don't ask for more pk attributes than we have
+	 * non-dropped columns
+	 */
+	nondropped_natts = get_nondropped_natts(relid);
+	if (pknumatts > nondropped_natts)
+		elog(ERROR, "number of primary key fields exceeds number of specified relation attributes");
+
 	src_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(3);
 	tgt_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(4);
 
@@ -2191,3 +2219,27 @@ dblink_security_check(PGconn *conn, const char *connstr)
 			PQfinish(conn);
 	}
 }
+
+static int
+get_nondropped_natts(Oid relid)
+{
+	int			nondropped_natts = 0;
+	TupleDesc	tupdesc;
+	Relation	rel;
+	int			natts;
+	int			i;
+
+	rel = relation_open(relid, AccessShareLock);
+	tupdesc = rel->rd_att;
+	natts = tupdesc->natts;
+
+	for (i = 0; i < natts; i++)
+	{
+		if (tupdesc->attrs[i]->attisdropped)
+			continue;
+		nondropped_natts++;
+	}
+
+	relation_close(rel, AccessShareLock);
+	return nondropped_natts;
+}
@@ -44,6 +44,9 @@ SELECT dblink_build_sql_insert('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
  INSERT INTO foo(f1,f2,f3) VALUES('99','xyz','{a0,b0,c0}')
 (1 row)
 
+-- too many pk fields, should fail
+SELECT dblink_build_sql_insert('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}','{"99", "xyz", "{za0,zb0,zc0}"}');
+ERROR:  number of primary key fields exceeds number of specified relation attributes
 -- build an update statement based on a local tuple,
 -- replacing the primary key values with new ones
 SELECT dblink_build_sql_update('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
@@ -52,13 +55,19 @@ SELECT dblink_build_sql_update('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
  UPDATE foo SET f1 = '99', f2 = 'xyz', f3 = '{a0,b0,c0}' WHERE f1 = '99' AND f2 = 'xyz'
 (1 row)
 
+-- too many pk fields, should fail
+SELECT dblink_build_sql_update('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}','{"99", "xyz", "{za0,zb0,zc0}"}');
+ERROR:  number of primary key fields exceeds number of specified relation attributes
 -- build a delete statement based on a local tuple,
 SELECT dblink_build_sql_delete('foo','1 2',2,'{"0", "a"}');
            dblink_build_sql_delete           
 ---------------------------------------------
  DELETE FROM foo WHERE f1 = '0' AND f2 = 'a'
 (1 row)
 
+-- too many pk fields, should fail
+SELECT dblink_build_sql_delete('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}');
+ERROR:  number of primary key fields exceeds number of specified relation attributes
 -- retest using a quoted and schema qualified table
 CREATE SCHEMA "MySchema";
 CREATE TABLE "MySchema"."Foo"(f1 int, f2 text, f3 text[], primary key (f1,f2));
 
@@ -36,13 +36,19 @@ FROM dblink_get_pkey('foo');
 -- build an insert statement based on a local tuple,
 -- replacing the primary key values with new ones
 SELECT dblink_build_sql_insert('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
+SELECT dblink_build_sql_insert('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}','{"99", "xyz", "{za0,zb0,zc0}"}');
 
 -- build an update statement based on a local tuple,
 -- replacing the primary key values with new ones
 SELECT dblink_build_sql_update('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
+-- too many pk fields, should fail
+SELECT dblink_build_sql_update('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}','{"99", "xyz", "{za0,zb0,zc0}"}');
 
 -- build a delete statement based on a local tuple,
 SELECT dblink_build_sql_delete('foo','1 2',2,'{"0", "a"}');
+-- too many pk fields, should fail
+SELECT dblink_build_sql_delete('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}');
 
 -- retest using a quoted and schema qualified table
 CREATE SCHEMA "MySchema";