Track BeginResultSet state to prevent formatter lifecycle violations

Co-authored-by: dlevy-msft-sql <[email protected]>

Author: Copilot

pkg/sqlcmd/sqlcmd.go

@@ -452,6 +452,7 @@ func (s *Sqlcmd) runQuery(query string) (int, error) {
 	var cols []*sql.ColumnType
 	results := true
 	first := true
+	beganResultSet := false
 	for qe == nil && results {
 		msg := retmsg.Message(ctx)
 		switch m := msg.(type) {
@@ -486,6 +487,7 @@ func (s *Sqlcmd) runQuery(query string) (int, error) {
 			}
 			if results {
 				first = true
+				beganResultSet = false
 			}
 		case sqlexp.MsgNext:
 			if first {
@@ -495,31 +497,35 @@ func (s *Sqlcmd) runQuery(query string) (int, error) {
 					retcode = -100
 					qe = s.handleError(&retcode, err)
 					s.Format.AddError(err)
-					// Skip processing this result set since we cannot get column types
-					continue
+					beganResultSet = false
+				} else {
+					s.Format.BeginResultSet(cols)
+					beganResultSet = true
 				}
-				s.Format.BeginResultSet(cols)
 			}
-			inresult := rows.Next()
-			for inresult {
-				col1 := s.Format.AddRow(rows)
-				inresult = rows.Next()
-				if !inresult {
-					if col1 == "" {
-						retcode = 0
-					} else if _, cerr := fmt.Sscanf(col1, "%d", &retcode); cerr != nil {
-						retcode = -102
+			// Only process rows if BeginResultSet was successfully called
+			if beganResultSet {
+				inresult := rows.Next()
+				for inresult {
+					col1 := s.Format.AddRow(rows)
+					inresult = rows.Next()
+					if !inresult {
+						if col1 == "" {
+							retcode = 0
+						} else if _, cerr := fmt.Sscanf(col1, "%d", &retcode); cerr != nil {
+							retcode = -102
+						}
 					}
 				}
-			}
-			if retcode != -102 {
-				if err = rows.Err(); err != nil {
-					retcode = -100
-					qe = s.handleError(&retcode, err)
-					s.Format.AddError(err)
+				if retcode != -102 {
+					if err = rows.Err(); err != nil {
+						retcode = -100
+						qe = s.handleError(&retcode, err)
+						s.Format.AddError(err)
+					}
 				}
+				s.Format.EndResultSet()
 			}
-			s.Format.EndResultSet()
 		}
 	}
 	s.Format.EndBatch()