Adding ClientOSHostname and Workload properties (#165)

Co-authored-by: sitingren <siting.ren@microfocus.com>

README.md

@@ -101,6 +101,7 @@ Currently supported query arguments are:
 | client_label   | Sets a label for the connection on the server. This value appears in the `client_label` column of the SESSIONS system table. | (default) vertica-sql-go-{version}-{pid}-{timestamp} |
 | autocommit     | Controls whether the connection automatically commits transactions. | 1 = (default) on <br>0 = off|
 | oauth_access_token | To authenticate via OAuth, provide an OAuth Access Token that authorizes a user to the database. | unspecified by default, if specified then *user* is optional |
+| workload | Sets workload property of the session, enabling use of workload routing | empty string by default. Valid values are workload names that already exist in a workload routing rule on the server. If a workload name that doesn't exist is entered, the server will reject it and it will be set to the default empty string |
 
 To ping the server and validate a connection (as the connection isn't necessarily created at that moment), simply call the *PingContext()* method.
 

connection.go

@@ -113,6 +113,7 @@ type connection struct {
 	serverTZOffset   string
 	dead             bool // used if a ROLLBACK severity error is encountered
 	sessMutex        sync.Mutex
+	workload         string
 }
 
 // Begin - Begin starts and returns a new transaction. (DEPRECATED)
@@ -260,6 +261,9 @@ func newConnection(connString string) (*connection, error) {
 		sslFlag = tlsModeNone
 	}
 
+	// Read Workload flag
+	result.workload = result.connURL.Query().Get("workload")
+
 	result.conn, err = result.establishSocketConnection()
 
 	if err != nil {
@@ -444,6 +448,7 @@ func (v *connection) handshake() error {
 		ClientPID:        v.clientPID,
 		Autocommit:       v.autocommit,
 		OAuthAccessToken: v.oauthaccesstoken,
+		Workload:         v.workload,
 	}
 
 	if err := v.sendMessage(msg); err != nil {

driver.go

@@ -47,7 +47,7 @@ type Driver struct{}
 const (
 	driverName      string = "vertica-sql-go"
 	driverVersion   string = "1.3.2"
-	protocolVersion uint32 = 0x0003000C
+	protocolVersion uint32 = 0x0003000F // 3.15
 )
 
 var driverLogger = logger.New("driver")

driver_test.go

@@ -178,14 +178,14 @@ func TestOAuthConnection(t *testing.T) {
 	assertNoErr(t, connDB.PingContext(ctx))
 
 	rows, err := connDB.QueryContext(ctx, "SELECT authentication_method FROM sessions WHERE session_id=(SELECT current_session())")
-    assertNoErr(t, err)
+	assertNoErr(t, err)
 	defer rows.Close()
 
 	var authMethod string
-    for rows.Next() {
-            assertNoErr(t, rows.Scan(&authMethod))
-			assertEqual(t, authMethod, "OAuth")
-    }
+	for rows.Next() {
+		assertNoErr(t, rows.Scan(&authMethod))
+		assertEqual(t, authMethod, "OAuth")
+	}
 }
 
 func TestTLSConfiguration(t *testing.T) {
@@ -1184,6 +1184,46 @@ func TestUnexpectedResult(t *testing.T) {
 	assertErr(t, err, "ERROR: whatever")
 }
 
+func TestWorkloadConnectionProperty(t *testing.T) {
+	connDB, err := sql.Open("vertica", myDBConnectString+"&workload=golangWorkload")
+	assertNoErr(t, err)
+
+	err = connDB.PingContext(ctx)
+	assertNoErr(t, err)
+	defer closeConnection(t, connDB)
+	rows, err := connDB.QueryContext(ctx, "SELECT contents FROM dc_client_server_messages "+
+		"WHERE session_id = current_session() "+
+		"AND message_type = '^+' "+
+		"AND contents LIKE '%workload%'")
+	assertNoErr(t, err)
+	defer rows.Close()
+
+	var workload string
+	for rows.Next() {
+		assertNoErr(t, rows.Scan(&workload))
+		assertEqual(t, workload, "workload: golangWorkload")
+	}
+}
+
+func TestClientOSHostnameProperty(t *testing.T) {
+	connDB := openConnection(t)
+	defer closeConnection(t, connDB)
+	rows, err := connDB.QueryContext(ctx, "SELECT client_os_hostname FROM current_session")
+	assertNoErr(t, err)
+	defer rows.Close()
+
+	var client_os_hostname = ""
+	hostname, err := os.Hostname()
+	if err == nil {
+		client_os_hostname = hostname
+	}
+	var server_side_client_os_hostname string
+	for rows.Next() {
+		assertNoErr(t, rows.Scan(&server_side_client_os_hostname))
+		assertEqual(t, server_side_client_os_hostname, client_os_hostname)
+	}
+}
+
 var verticaUserName = flag.String("user", "dbadmin", "the user name to connect to Vertica")
 var verticaPassword = flag.String("password", os.Getenv("VERTICA_TEST_PASSWORD"), "Vertica password for this user")
 var verticaHostPort = flag.String("locator", "localhost:5433", "Vertica's host and port")
@@ -1271,7 +1311,7 @@ func init() {
 	otherConnectString = "vertica://TestGuy:TestGuyPass@" + *verticaHostPort + "/?tlsmode=" + *tlsMode
 	badConnectString = "vertica://TestGuy:TestGuyBadPass@" + *verticaHostPort + "/?tlsmode=" + *tlsMode
 	failoverConnectString = "vertica://" + *verticaUserName + ":" + *verticaPassword + "@badHost" + "?backup_server_node=abc.com:100000," + *verticaHostPort + ",localhost:port"
-	oauthConnectString = "vertica://" + *verticaUserName + "@" + *verticaHostPort +  "/?oauth_access_token=" + *oauthAccessToken + "&tlsmode=" + *tlsMode
+	oauthConnectString = "vertica://" + *verticaUserName + "@" + *verticaHostPort + "/?oauth_access_token=" + *oauthAccessToken + "&tlsmode=" + *tlsMode
 
 	ctx = context.Background()
 }

msgs/festartupmsg.go

@@ -34,6 +34,7 @@ package msgs
 
 import (
 	"fmt"
+	"os"
 	"os/user"
 
 	"github.com/elastic/go-sysinfo"
@@ -52,6 +53,8 @@ type FEStartupMsg struct {
 	OSUsername       string
 	Autocommit       string
 	OAuthAccessToken string
+	ClientOSHostname string
+	Workload         string
 }
 
 // Flatten docs
@@ -70,6 +73,12 @@ func (m *FEStartupMsg) Flatten() ([]byte, byte) {
 		m.OSUsername = currentUser.Username
 	}
 
+	m.ClientOSHostname = ""
+	hostname, err := os.Hostname()
+	if err == nil {
+		m.ClientOSHostname = hostname
+	}
+
 	buf := newMsgBuffer()
 	const fixedProtocolVersion uint32 = 0x00030005
 	buf.appendUint32(fixedProtocolVersion)
@@ -97,14 +106,17 @@ func (m *FEStartupMsg) Flatten() ([]byte, byte) {
 	buf.appendLabeledString("client_os", m.ClientOS)
 	buf.appendLabeledString("client_os_user_name", m.OSUsername)
 	buf.appendLabeledString("autocommit", m.Autocommit)
+	buf.appendLabeledString("protocol_compat", "VER")
+	buf.appendLabeledString("client_os_hostname", m.ClientOSHostname)
+	buf.appendLabeledString("workload", m.Workload)
 	buf.appendBytes([]byte{0})
 
 	return buf.bytes(), 0
 }
 
 func (m *FEStartupMsg) String() string {
 	return fmt.Sprintf(
-		"Startup (packet): ProtocolVersion:%08X, DriverName='%s', DriverVersion='%s', UserName='%s', Database='%s', SessionID='%s', ClientPID=%d, ClientOS='%s', ClientOSUserName='%s', Autocommit='%s', OAuthAccessToken=<length:%d>",
+		"Startup (packet): ProtocolVersion:%08X, DriverName='%s', DriverVersion='%s', UserName='%s', Database='%s', SessionID='%s', ClientPID=%d, ClientOS='%s', ClientOSUserName='%s', ClientOSHostname='%s', Autocommit='%s', OAuthAccessToken=<length:%d>, Workload='%s'",
 		m.ProtocolVersion,
 		m.DriverName,
 		m.DriverVersion,
@@ -114,6 +126,8 @@ func (m *FEStartupMsg) String() string {
 		m.ClientPID,
 		m.ClientOS,
 		m.OSUsername,
+		m.ClientOSHostname,
 		m.Autocommit,
-		len(m.OAuthAccessToken))
+		len(m.OAuthAccessToken),
+		m.Workload)
 }