Fix for panic in rows.Next() when stored procedures emit NOTICE messages (#188)

Author: sivaalamp

driver_test.go

@@ -1660,6 +1660,128 @@ func TestInvalidEmailParseStatement(t *testing.T) {
 		assertEqual(t, all_roles, role)
 	}
 }
+
+// TestStoredProcedureWithNotice verifies that rows.Next() does not panic when a
+// stored procedure raises one or more RAISE NOTICE messages before returning its
+// result set.
+//
+// Root cause: Vertica sends a correct RowDescription (2 cols: a, b) for the CALL
+// result, then emits NoticeResponse messages, then sends a spurious second
+// RowDescription with fewer columns (a PLvSQL protocol artefact).  The driver was
+// blindly adopting the second RowDescription, so r.columnDefs ended up with 1
+// column while the DataRow still contained 2 fields — causing an index-out-of-range
+// panic in rows.Next().
+//
+// Fix: runSimpleStatement and collectResults now ignore any RowDescription that
+// arrives after DataRows have already been buffered for the current result set.
+func TestStoredProcedureWithNotice(t *testing.T) {
+	connDB := openConnection(t)
+	defer closeConnection(t, connDB)
+
+	// Create the procedure; drop it afterward.
+	_, err := connDB.ExecContext(ctx, `
+		CREATE OR REPLACE PROCEDURE test_notice_proc(INOUT a INT, INOUT b VARCHAR(64))
+		LANGUAGE PLvSQL AS $$
+		BEGIN
+			RAISE NOTICE 'Value of a: %', a;
+			RAISE NOTICE 'Value of b: %', b;
+		END;
+		$$`)
+	assertNoErr(t, err)
+
+	defer connDB.ExecContext(ctx, "DROP PROCEDURE IF EXISTS test_notice_proc(INT, VARCHAR)")
+
+	// CALL must not panic.  We wrap in a recover so that a panic becomes a
+	// test failure with a meaningful message rather than crashing the whole
+	// test binary.
+	func() {
+		defer func() {
+			if r := recover(); r != nil {
+				t.Fatalf("rows.Next() panicked when stored procedure emitted NOTICE messages: %v", r)
+			}
+		}()
+
+		rows, err := connDB.QueryContext(ctx, "CALL test_notice_proc(10, 'hello')")
+		assertNoErr(t, err)
+		defer rows.Close()
+
+		// Iterate over all result sets (NOTICE messages may produce extra sets).
+		for {
+			cols, err := rows.Columns()
+			assertNoErr(t, err)
+
+			valuePtrs := make([]interface{}, len(cols))
+			values := make([]interface{}, len(cols))
+			for i := range values {
+				valuePtrs[i] = &values[i]
+			}
+
+			for rows.Next() {
+				assertNoErr(t, rows.Scan(valuePtrs...))
+				testLogger.Debug("TestStoredProcedureWithNotice row: %v", values)
+			}
+
+			if !rows.NextResultSet() {
+				break
+			}
+		}
+	}()
+}
+
+// TestStoredProcedureWithNoticeSimpleQuery is the same scenario as
+// TestStoredProcedureWithNotice but forces the simple query protocol path
+// (use_prepared_statements=0) to verify correct behaviour on both code paths.
+func TestStoredProcedureWithNoticeSimpleQuery(t *testing.T) {
+	simpleConnStr := strings.Replace(myDBConnectString, "use_prepared_statements=1", "use_prepared_statements=0", 1)
+	connDB, err := sql.Open("vertica", simpleConnStr)
+	assertNoErr(t, err)
+	assertNoErr(t, connDB.PingContext(ctx))
+	defer closeConnection(t, connDB)
+
+	_, err = connDB.ExecContext(ctx, `
+		CREATE OR REPLACE PROCEDURE test_notice_proc_simple(INOUT a INT, INOUT b VARCHAR(64))
+		LANGUAGE PLvSQL AS $$
+		BEGIN
+			RAISE NOTICE 'Value of a: %', a;
+			RAISE NOTICE 'Value of b: %', b;
+		END;
+		$$`)
+	assertNoErr(t, err)
+	defer connDB.ExecContext(ctx, "DROP PROCEDURE IF EXISTS test_notice_proc_simple(INT, VARCHAR)")
+
+	func() {
+		defer func() {
+			if r := recover(); r != nil {
+				t.Fatalf("rows.Next() panicked on simple query path: %v", r)
+			}
+		}()
+
+		rows, err := connDB.QueryContext(ctx, "CALL test_notice_proc_simple(10, 'hello')")
+		assertNoErr(t, err)
+		defer rows.Close()
+
+		for {
+			cols, err := rows.Columns()
+			assertNoErr(t, err)
+
+			valuePtrs := make([]interface{}, len(cols))
+			values := make([]interface{}, len(cols))
+			for i := range values {
+				valuePtrs[i] = &values[i]
+			}
+
+			for rows.Next() {
+				assertNoErr(t, rows.Scan(valuePtrs...))
+				testLogger.Debug("TestStoredProcedureWithNoticeSimpleQuery row: %v", values)
+			}
+
+			if !rows.NextResultSet() {
+				break
+			}
+		}
+	}()
+}
+
 func init() {
 	// One or both lines below are necessary depending on your go version
 	testing.Init()

rows.go

@@ -352,3 +352,27 @@ func newEmptyRows() *rows {
 	be := &msgs.BERowDescMsg{Columns: cdf}
 	return newRows(context.Background(), be, "")
 }
+
+// expandColumnDefs grows r.columnDefs to cover at least numCols columns.
+// It is a last-resort fallback for CALL statements where Vertica's Describe
+// phase returns an incomplete RowDescription and no execution-time
+// RowDescription is sent to correct it. Synthetic entries use OID 0 so that
+// rows.Next() returns the raw wire bytes as a string (the default case),
+// making the behaviour explicit rather than falsely claiming a specific type.
+// The caller invokes this whenever a DataRow wider than columnDefs is seen,
+// not only on the first row.
+//
+// Ordering guarantee: this function is called inside collectResults(), which
+// fully drains the wire and buffers all rows before returning the *rows object
+// to database/sql. Because database/sql only calls Columns() after receiving
+// that object, all expansions are complete before Columns() executes.
+// This invariant would break in a streaming/lazy-row model where *rows is
+// handed to the caller before all DataRows have been received.
+func (r *rows) expandColumnDefs(numCols uint16) {
+	for uint16(len(r.columnDefs.Columns)) < numCols {
+		r.columnDefs.Columns = append(r.columnDefs.Columns, &msgs.BERowDescColumnDef{
+			FieldName:   fmt.Sprintf("unknown_col%d", len(r.columnDefs.Columns)),
+			DataTypeOID: 0, // unknown — falls to the default string case in rows.Next()
+		})
+	}
+}

stmt.go

@@ -595,13 +595,32 @@ func (s *stmt) collectResults(ctx context.Context) (*rows, error) {
 
 		switch msg := bMsg.(type) {
 		case *msgs.BEDataRowMsg:
+			// Vertica's Describe phase under-reports the column count for CALL
+			// statements that emit RAISE NOTICE: a DataRow at execution time may
+			// contain more fields than columnDefs describes. Expand columnDefs
+			// whenever a wider DataRow is seen so that Columns() always returns
+			// the correct width.
+			//
+			// This is safe because collectResults() is synchronous and fully
+			// buffers all rows before returning *rows to the caller. database/sql
+			// calls Columns() only after receiving that object, so all expansions
+			// are complete by the time the column list is observed.
+			if uint16(len(rows.columnDefs.Columns)) < msg.Columns().NumCols {
+				rows.expandColumnDefs(msg.Columns().NumCols)
+			}
 			err = rows.addRow(msg)
 			if err != nil {
 				return rows, err
 			}
 		case *msgs.BERowDescMsg:
-			s.lastRowDesc = msg
-			rows = newRows(ctx, s.lastRowDesc, s.conn.serverTZOffset)
+			// An execution-time RowDescription may arrive before any DataRows
+			// and can carry a more complete schema than the Describe-phase one.
+			// Only adopt it if it has at least as many columns, to prevent a
+			// truncated description from silently replacing a wider one.
+			if rows.resultData.Peek() == nil && len(msg.Columns) >= len(rows.columnDefs.Columns) {
+				s.lastRowDesc = msg
+				rows = newRows(ctx, s.lastRowDesc, s.conn.serverTZOffset)
+			}
 		case *msgs.BEErrorMsg:
 			s.conn.sync()
 			return newEmptyRows(), s.evaluateErrorMsg(msg)