Set correct identity before run_statement

Signed-off-by: mohitkhullar <mkhullar1@bloomberg.net>

Author: mohitkhullar

db/fdb_fend.c

@@ -4293,6 +4293,15 @@ void fdb_free_tran(sqlclntstate *clnt, fdb_tran_t *tran)
 
 extern char gbl_dbname[];
 
+void fdb_client_set_identityBlob(sqlclntstate *clnt, cdb2_hndl_tp *hndl)
+{
+    extern void *(*externalComdb2getAuthIdBlob)(void *ID);
+    if (gbl_fdb_auth_enabled && externalComdb2getAuthIdBlob &&
+        ((clnt->authdata = get_authdata(clnt)) != NULL)) {
+        cdb2_setIdentityBlob(hndl, externalComdb2getAuthIdBlob(clnt->authdata));
+    }
+}
+
 int fdb_trans_commit(sqlclntstate *clnt, enum trans_clntcomm sideeffects)
 {
     fdb_distributed_tran_t *dtran = clnt->dbtran.dtran;
@@ -4340,6 +4349,7 @@ int fdb_trans_commit(sqlclntstate *clnt, enum trans_clntcomm sideeffects)
             if (tran->nwrites) {
                 /* handle is only created upon first remote write to this fdb */
                 assert(tran->fcon.hndl);
+                fdb_client_set_identityBlob(clnt, tran->fcon.hndl);
                 rc = cdb2_run_statement(tran->fcon.hndl, "commit");
             } else {
                 rc = 0;
@@ -4461,6 +4471,7 @@ int fdb_trans_rollback(sqlclntstate *clnt)
             if (tran->nwrites) {
                 /* handle is only created upon first remote write to this fdb */
                 assert(tran->fcon.hndl);
+                fdb_client_set_identityBlob(clnt, tran->fcon.hndl);
                 rc = cdb2_run_statement(tran->fcon.hndl, "rollback");
             } else {
                 rc = 0;
@@ -5748,12 +5759,7 @@ static int _fdb_client_set_options(sqlclntstate *clnt,
     if (clnt->prepare_only) {
         SET_STR("PREPARE_ONLY", "ON");
     }
-
-    extern void *(*externalComdb2getAuthIdBlob)(void *ID);
-    if (gbl_fdb_auth_enabled && externalComdb2getAuthIdBlob &&
-        ((clnt->authdata = get_authdata(clnt)) != NULL)) {
-        cdb2_setIdentityBlob(hndl, externalComdb2getAuthIdBlob(clnt->authdata));
-    }
+    fdb_client_set_identityBlob(clnt, hndl);
 
     return 0;
 }

db/fdb_push.c

@@ -308,6 +308,8 @@ static int forward_extra_set_commands(cdb2_hndl_tp *hndl, int n_sets, const char
     return 0;
 }
 
+void fdb_client_set_identityBlob(sqlclntstate *clnt, cdb2_hndl_tp *hndl);
+
 static cdb2_hndl_tp *_hndl_open_int(sqlclntstate *clnt, const char *class,
                                     int flags, struct errstat *err,
                                     int n_sets, const char **sets)
@@ -354,9 +356,7 @@ static cdb2_hndl_tp *_hndl_open_int(sqlclntstate *clnt, const char *class,
         }
     }
 
-    if (gbl_fdb_auth_enabled && externalComdb2getAuthIdBlob &&
-         ((clnt->authdata = get_authdata(clnt)) != NULL))
-        cdb2_setIdentityBlob(hndl, externalComdb2getAuthIdBlob(clnt->authdata));
+    fdb_client_set_identityBlob(clnt, hndl);
 
     return hndl;
 }
@@ -615,6 +615,8 @@ int handle_fdb_push_write(sqlclntstate *clnt, struct errstat *err,
     }
     hndl = tran->fcon.hndl;
 
+    fdb_client_set_identityBlob(clnt, hndl);
+
     /* run the statement */
     rc = _run_statement(clnt, hndl, err);
     if (rc != CDB2_OK) {