@@ -86,6 +86,12 @@ type Client struct {
8686
8787 // for managing one-time messaging
8888 oneTimeMessages map [string ]struct {}
89+
90+ // sessionStr represents a session cookie to use, if non-nil
91+ sessionStr string
92+
93+ // for generic locking
94+ lockMux sync.Mutex
8995}
9096
9197var (
97103const (
98104 // LimiterRefreshNanos is used to update table limits once every 10 minutes
99105 LimiterRefreshNanos int64 = 600 * 1000 * 1000 * 1000
106+ // SessionCookieField is used to check for persistent session cookies
107+ SessionCookieField string = "session="
100108)
101109
102110// NewClient creates a Client instance with the specified Config.
@@ -794,9 +802,9 @@ func (c *Client) nextRequestID() int32 {
794802// processRequest processes the specified request before it is sent to server.
795803// This method applies default configurations such as timeout and consistency
796804// values for the request if they are not specified for the request.
797- func (c * Client ) processRequest (req Request ) (data []byte , err error ) {
805+ func (c * Client ) processRequest (req Request ) (data []byte , serialVerUsed int16 , err error ) {
798806 if req == nil {
799- return nil , errNilRequest
807+ return nil , 0 , errNilRequest
800808 }
801809
802810 // Set default values for the request with the global request configurations
@@ -806,17 +814,17 @@ func (c *Client) processRequest(req Request) (data []byte, err error) {
806814
807815 // Validates the request, returns immediately if validation fails.
808816 if err = req .validate (); err != nil {
809- return nil , err
817+ return nil , 0 , err
810818 }
811819
812- data , err = c .serializeRequest (req )
820+ data , serialVerUsed , err = c .serializeRequest (req )
813821 if err != nil || ! c .isCloud {
814822 return
815823 }
816824
817825 // check request size for cloud
818826 if err = checkRequestSizeLimit (req , len (data )); err != nil {
819- return nil , err
827+ return nil , 0 , err
820828 }
821829
822830 return
@@ -831,15 +839,15 @@ func (c *Client) execute(req Request) (Result, error) {
831839}
832840
833841func (c * Client ) executeWithContext (ctx context.Context , req Request ) (Result , error ) {
834- data , err := c .processRequest (req )
842+ data , serialVerUsed , err := c .processRequest (req )
835843 if err != nil {
836844 return nil , err
837845 }
838846
839- return c .doExecute (ctx , req , data )
847+ return c .doExecute (ctx , req , data , serialVerUsed )
840848}
841849
842- func (c * Client ) doExecute (ctx context.Context , req Request , data []byte ) (result Result , err error ) {
850+ func (c * Client ) doExecute (ctx context.Context , req Request , data []byte , serialVerUsed int16 ) (result Result , err error ) {
843851 if req == nil {
844852 return nil , errNilRequest
845853 }
@@ -962,11 +970,11 @@ func (c *Client) doExecute(ctx context.Context, req Request, data []byte) (resul
962970 }
963971
964972 if nosqlerr .Is (err , nosqlerr .UnsupportedProtocol ) {
965- if c .decrementSerialVersion () == false {
973+ if c .decrementSerialVersion (serialVerUsed ) == false {
966974 return nil , err
967975 }
968976 // if serial version mismatch, we must re-serialize the request
969- data , err = c .serializeRequest (req )
977+ data , serialVerUsed , err = c .serializeRequest (req )
970978 if err != nil {
971979 return nil , err
972980 }
@@ -1048,14 +1056,19 @@ func (c *Client) doExecute(ctx context.Context, req Request, data []byte) (resul
10481056 httpReq .Header .Set ("Authorization" , authStr )
10491057 }
10501058
1059+ // Allow for session persistence, if available
1060+ if c .sessionStr != "" {
1061+ httpReq .Header .Set ("Cookie" , c .sessionStr )
1062+ }
1063+
10511064 err = c .signHTTPRequest (httpReq )
10521065 if err != nil {
10531066 return nil , err
10541067 }
10551068
10561069 // warn if using features not implemented at the connected server
10571070 // currently cloud does not support Durability
1058- if c . serialVersion < 3 || c .isCloud {
1071+ if serialVerUsed < 3 || c .isCloud {
10591072 needMsg := false
10601073 if pReq , ok := req .(* PutRequest ); ok && pReq .Durability .IsSet () {
10611074 needMsg = true
@@ -1073,7 +1086,7 @@ func (c *Client) doExecute(ctx context.Context, req Request, data []byte) (resul
10731086 }
10741087
10751088 // OnDemand is not available in V2
1076- if c . serialVersion < 3 {
1089+ if serialVerUsed < 3 {
10771090 if tReq , ok := req .(* TableRequest ); ok && tReq .TableLimits != nil {
10781091 if tReq .TableLimits .CapacityMode == types .OnDemand {
10791092 c .oneTimeMessage ("The requested feature is not supported " +
@@ -1334,17 +1347,18 @@ func (c *Client) signHTTPRequest(httpReq *http.Request) error {
13341347// serializeRequest serializes the specified request into a slice of bytes that
13351348// will be sent to the server. The serial version is always written followed by
13361349// the actual request payload.
1337- func (c * Client ) serializeRequest (req Request ) (data []byte , err error ) {
1350+ func (c * Client ) serializeRequest (req Request ) (data []byte , serialVerUsed int16 , err error ) {
13381351 wr := binary .NewWriter ()
1339- if _ , err = wr .WriteSerialVersion (c .serialVersion ); err != nil {
1340- return
1352+ serialVerUsed = c .serialVersion
1353+ if _ , err = wr .WriteSerialVersion (serialVerUsed ); err != nil {
1354+ return nil , 0 , err
13411355 }
13421356
1343- if err = req .serialize (wr , c . serialVersion ); err != nil {
1344- return
1357+ if err = req .serialize (wr , serialVerUsed ); err != nil {
1358+ return nil , 0 , err
13451359 }
13461360
1347- return wr .Bytes (), nil
1361+ return wr .Bytes (), serialVerUsed , nil
13481362}
13491363
13501364// processResponse processes the http response returned from server.
@@ -1360,6 +1374,7 @@ func (c *Client) processResponse(httpResp *http.Response, req Request) (Result,
13601374 }
13611375
13621376 if httpResp .StatusCode == http .StatusOK {
1377+ c .setSessionCookie (httpResp .Header )
13631378 return c .processOKResponse (data , req )
13641379 }
13651380
@@ -1402,6 +1417,29 @@ func (c *Client) processOKResponse(data []byte, req Request) (Result, error) {
14021417 return nil , wrapResponseErrors (int (code ), msg )
14031418}
14041419
1420+ // setSessionCookie sets a persistent session cookie value to use for
1421+ // following requests, if present in the response header.
1422+ func (c * Client ) setSessionCookie (header http.Header ) {
1423+ if header == nil {
1424+ return
1425+ }
1426+ // NOTE: this code assumes there will always be at most
1427+ // one Set-Cookie header in the response. If the load balancer
1428+ // settings change, or the proxy changes to add Set-Cookie
1429+ // headers, this code may need to be changed to look for
1430+ // multiple Set-Cookie headers.
1431+ v := header .Get ("Set-Cookie" )
1432+ if strings .HasPrefix (v , SessionCookieField ) == false {
1433+ return
1434+ }
1435+ c .lockMux .Lock ()
1436+ defer c .lockMux .Unlock ()
1437+ c .sessionStr = strings .Split (v , ";" )[0 ]
1438+ c .logger .LogWithFn (logger .Fine , func () string {
1439+ return fmt .Sprintf ("Set session cookie to \" %s\" " , c .sessionStr )
1440+ })
1441+ }
1442+
14051443// processNotOKResponse processes the http response whose status code is not 200.
14061444func (c * Client ) processNotOKResponse (data []byte , statusCode int ) error {
14071445 if statusCode == http .StatusBadRequest && len (data ) > 0 {
@@ -1502,7 +1540,12 @@ func (c *Client) VerifyConnection() error {
15021540// decrementSerialVersion attempts to reduce the serial version used for
15031541// communicating with the server. If the version is already at its lowest
15041542// value, it will not be decremented and false will be returned.
1505- func (c * Client ) decrementSerialVersion () bool {
1543+ func (c * Client ) decrementSerialVersion (serialVerUsed int16 ) bool {
1544+ c .lockMux .Lock ()
1545+ defer c .lockMux .Unlock ()
1546+ if c .serialVersion != serialVerUsed {
1547+ return true
1548+ }
15061549 if c .serialVersion > 2 {
15071550 c .serialVersion --
15081551 return true
0 commit comments