From eff882af152be9731c0a4db21e3c117d04564f80 Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Wed, 27 Oct 2021 18:40:55 -0700
Subject: [PATCH] sql: fix OIDs in RowDescription in some cases

Release note (bug fix): Previously, CockroachDB could not set the
`TableOID` and `TableAttributeNumber` attributes of `RowDescription`
message of pgwire protocol in some cases (these values would be left as
0), and this is now fixed.
---
 pkg/sql/exec_factory_util.go                  | 12 ++++----
 pkg/sql/opt_exec_factory.go                   |  7 +----
 .../pgwire/testdata/pgtest/row_description    | 29 ++++++++++++++++---
 3 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/pkg/sql/exec_factory_util.go b/pkg/sql/exec_factory_util.go
index 48d131832c0e..e7e49cc85615 100644
--- a/pkg/sql/exec_factory_util.go
+++ b/pkg/sql/exec_factory_util.go
@@ -123,11 +123,11 @@ func makeScanColumnsConfig(table cat.Table, cols exec.TableColumnOrdinalSet) sca
 }
 
 // getResultColumnsForSimpleProject populates result columns for a simple
-// projection. It supports two configurations:
+// projection. inputCols must be non-nil and contain the result columns before
+// the projection has been applied. It supports two configurations:
 // 1. colNames and resultTypes are non-nil. resultTypes indicates the updated
 //    types (after the projection has been applied)
-// 2. if colNames is nil, then inputCols must be non-nil (which are the result
-//    columns before the projection has been applied).
+// 2. colNames is nil.
 func getResultColumnsForSimpleProject(
 	cols []exec.NodeColumnOrdinal,
 	colNames []string,
@@ -143,8 +143,10 @@ func getResultColumnsForSimpleProject(
 			resultCols[i].Hidden = false
 		} else {
 			resultCols[i] = colinfo.ResultColumn{
-				Name: colNames[i],
-				Typ:  resultTypes[i],
+				Name:           colNames[i],
+				Typ:            resultTypes[i],
+				TableID:        inputCols[col].TableID,
+				PGAttributeNum: inputCols[col].PGAttributeNum,
 			}
 		}
 	}
diff --git a/pkg/sql/opt_exec_factory.go b/pkg/sql/opt_exec_factory.go
index 31d939e90220..5b7a23f43a65 100644
--- a/pkg/sql/opt_exec_factory.go
+++ b/pkg/sql/opt_exec_factory.go
@@ -247,11 +247,6 @@ func constructSimpleProjectForPlanNode(
 		r.reqOrdering = ReqOrdering(reqOrdering)
 		return r, nil
 	}
-	var inputCols colinfo.ResultColumns
-	if colNames == nil {
-		// We will need the names of the input columns.
-		inputCols = planColumns(n.(planNode))
-	}
 
 	var rb renderBuilder
 	rb.init(n, reqOrdering)
@@ -268,7 +263,7 @@ func constructSimpleProjectForPlanNode(
 			resultTypes[i] = exprs[i].ResolvedType()
 		}
 	}
-	resultCols := getResultColumnsForSimpleProject(cols, colNames, resultTypes, inputCols)
+	resultCols := getResultColumnsForSimpleProject(cols, colNames, resultTypes, planColumns(n.(planNode)))
 	rb.setOutput(exprs, resultCols)
 	return rb.res, nil
 }
diff --git a/pkg/sql/pgwire/testdata/pgtest/row_description b/pkg/sql/pgwire/testdata/pgtest/row_description
index defee35a1477..418ab6b5161e 100644
--- a/pkg/sql/pgwire/testdata/pgtest/row_description
+++ b/pkg/sql/pgwire/testdata/pgtest/row_description
@@ -123,12 +123,33 @@ RowDescription
 ----
 {"Type":"RowDescription","Fields":[{"Name":"v1","TableOID":0,"TableAttributeNumber":1,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0},{"Name":"v2","TableOID":0,"TableAttributeNumber":2,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
 
-# TODO(yuzefovich): we should not be ignoring table OIDs here, but currently we
-# have a bug there (#71891).
-until crdb_only ignore_table_oids
+until crdb_only
+RowDescription
+----
+{"Type":"RowDescription","Fields":[{"Name":"v1","TableOID":52,"TableAttributeNumber":1,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0},{"Name":"v2","TableOID":53,"TableAttributeNumber":2,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
+
+until ignore_table_oids
+ReadyForQuery
+----
+{"Type":"DataRow","Values":[{"text":"1"},{"text":"1"}]}
+{"Type":"CommandComplete","CommandTag":"SELECT 1"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
+# Regression test for not setting OIDs in some cases (#71891).
+send
+Query {"String": "SELECT a, tab1_a FROM tab2 INNER MERGE JOIN tab1 ON a = tab1_a WHERE a = 1"}
+----
+
+# With postgres we don't control the table OID.
+until ignore_table_oids noncrdb_only
+RowDescription
+----
+{"Type":"RowDescription","Fields":[{"Name":"a","TableOID":0,"TableAttributeNumber":1,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0},{"Name":"tab1_a","TableOID":0,"TableAttributeNumber":2,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
+
+until crdb_only
 RowDescription
 ----
-{"Type":"RowDescription","Fields":[{"Name":"v1","TableOID":0,"TableAttributeNumber":0,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0},{"Name":"v2","TableOID":0,"TableAttributeNumber":0,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
+{"Type":"RowDescription","Fields":[{"Name":"a","TableOID":52,"TableAttributeNumber":1,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0},{"Name":"tab1_a","TableOID":53,"TableAttributeNumber":2,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
 
 until ignore_table_oids
 ReadyForQuery