From 6bcffe14cd06dcdf9e9ceb8e5a0a4b21f2f2c5be Mon Sep 17 00:00:00 2001
From: testisnullus <danil.savchenko25@gmail.com>
Date: Mon, 22 Jul 2024 15:34:30 +0300
Subject: [PATCH] Added uint32 flags for QUERY and BATCH operations for native
 protocol v5

---
 frame.go | 101 ++++++++++++++++++++++++++++++++-----------------------
 1 file changed, 59 insertions(+), 42 deletions(-)

diff --git a/frame.go b/frame.go
index 44be7879d..aa1846093 100644
--- a/frame.go
+++ b/frame.go
@@ -150,14 +150,14 @@ const (
 	flagNoMetaData      int = 0x04
 
 	// query flags
-	flagValues                byte = 0x01
-	flagSkipMetaData          byte = 0x02
-	flagPageSize              byte = 0x04
-	flagWithPagingState       byte = 0x08
-	flagWithSerialConsistency byte = 0x10
-	flagDefaultTimestamp      byte = 0x20
-	flagWithNameValues        byte = 0x40
-	flagWithKeyspace          byte = 0x80
+	flagValues                uint32 = 0x01
+	flagSkipMetaData                 = 0x02
+	flagPageSize                     = 0x04
+	flagWithPagingState              = 0x08
+	flagWithSerialConsistency        = 0x10
+	flagDefaultTimestamp             = 0x20
+	flagWithNameValues               = 0x40
+	flagWithKeyspace                 = 0x80
 
 	// prepare flags
 	flagWithPreparedKeyspace uint32 = 0x01
@@ -1452,49 +1452,66 @@ func (f *framer) writeQueryParams(opts *queryParams) {
 		return
 	}
 
-	var flags byte
-	if len(opts.values) > 0 {
-		flags |= flagValues
-	}
-	if opts.skipMeta {
-		flags |= flagSkipMetaData
-	}
-	if opts.pageSize > 0 {
-		flags |= flagPageSize
-	}
-	if len(opts.pagingState) > 0 {
-		flags |= flagWithPagingState
-	}
-	if opts.serialConsistency > 0 {
-		flags |= flagWithSerialConsistency
-	}
-
+	var flagsByte byte
+	var flagsInt uint32
 	names := false
 
-	// protoV3 specific things
-	if f.proto > protoVersion2 {
+	switch {
+	case f.proto < protoVersion5:
+		if len(opts.values) > 0 {
+			flagsByte |= byte(flagValues)
+		}
+		if opts.skipMeta {
+			flagsByte |= byte(flagSkipMetaData)
+		}
+		if opts.pageSize > 0 {
+			flagsByte |= byte(flagPageSize)
+		}
+		if len(opts.pagingState) > 0 {
+			flagsByte |= byte(flagPageSize)
+		}
+		if opts.serialConsistency > 0 {
+			flagsByte |= byte(flagWithSerialConsistency)
+		}
+		if f.proto > protoVersion2 {
+			if opts.defaultTimestamp {
+				flagsByte |= byte(flagDefaultTimestamp)
+			}
+			if len(opts.values) > 0 && opts.values[0].name != "" {
+				flagsByte |= byte(flagWithNameValues)
+				names = true
+			}
+		}
+		f.writeByte(flagsByte)
+
+	case f.proto >= protoVersion5:
+		if len(opts.values) > 0 {
+			flagsInt |= flagValues
+		}
+		if opts.skipMeta {
+			flagsInt |= flagSkipMetaData
+		}
+		if opts.pageSize > 0 {
+			flagsInt |= flagPageSize
+		}
+		if len(opts.pagingState) > 0 {
+			flagsInt |= flagWithPagingState
+		}
+		if opts.serialConsistency > 0 {
+			flagsInt |= flagWithSerialConsistency
+		}
 		if opts.defaultTimestamp {
-			flags |= flagDefaultTimestamp
+			flagsInt |= flagDefaultTimestamp
 		}
 
 		if len(opts.values) > 0 && opts.values[0].name != "" {
-			flags |= flagWithNameValues
+			flagsInt |= flagWithNameValues
 			names = true
 		}
-	}
-
-	if opts.keyspace != "" {
-		if f.proto > protoVersion4 {
-			flags |= flagWithKeyspace
-		} else {
-			panic(fmt.Errorf("the keyspace can only be set with protocol 5 or higher"))
+		if opts.keyspace != "" {
+			flagsInt |= flagWithKeyspace
 		}
-	}
-
-	if f.proto > protoVersion4 {
-		f.writeUint(uint32(flags))
-	} else {
-		f.writeByte(flags)
+		f.writeUint(flagsInt)
 	}
 
 	if n := len(opts.values); n > 0 {