@@ -16,14 +16,16 @@ import (
1616 "github.com/pkg/errors"
1717)
1818
19- // this is used to generate test data. Developer should change this manually
19+ // RecordResults is used to generate test data. Developer should change this manually
2020var RecordResults bool
2121var resultIndex int
2222
2323type ThriftServiceClient struct {
2424 * cli_service.TCLIServiceClient
2525}
2626
27+ // OpenSession is a wrapper around the thrift operation OpenSession
28+ // If RecordResults is true, the results will be marshalled to JSON format and written to OpenSession<index>.json
2729func (tsc * ThriftServiceClient ) OpenSession (ctx context.Context , req * cli_service.TOpenSessionReq ) (* cli_service.TOpenSessionResp , error ) {
2830 msg , start := logger .Track ("OpenSession" )
2931 resp , err := tsc .TCLIServiceClient .OpenSession (ctx , req )
@@ -40,6 +42,8 @@ func (tsc *ThriftServiceClient) OpenSession(ctx context.Context, req *cli_servic
4042 return resp , CheckStatus (resp )
4143}
4244
45+ // CloseSession is a wrapper around the thrift operation CloseSession
46+ // If RecordResults is true, the results will be marshalled to JSON format and written to CloseSession<index>.json
4347func (tsc * ThriftServiceClient ) CloseSession (ctx context.Context , req * cli_service.TCloseSessionReq ) (* cli_service.TCloseSessionResp , error ) {
4448 log := logger .WithContext (driverctx .ConnIdFromContext (ctx ), driverctx .CorrelationIdFromContext (ctx ), "" )
4549 defer log .Duration (logger .Track ("CloseSession" ))
@@ -55,6 +59,8 @@ func (tsc *ThriftServiceClient) CloseSession(ctx context.Context, req *cli_servi
5559 return resp , CheckStatus (resp )
5660}
5761
62+ // FetchResults is a wrapper around the thrift operation FetchResults
63+ // If RecordResults is true, the results will be marshalled to JSON format and written to FetchResults<index>.json
5864func (tsc * ThriftServiceClient ) FetchResults (ctx context.Context , req * cli_service.TFetchResultsReq ) (* cli_service.TFetchResultsResp , error ) {
5965 log := logger .WithContext (driverctx .ConnIdFromContext (ctx ), driverctx .CorrelationIdFromContext (ctx ), SprintGuid (req .OperationHandle .OperationId .GUID ))
6066 defer log .Duration (logger .Track ("FetchResults" ))
@@ -70,6 +76,8 @@ func (tsc *ThriftServiceClient) FetchResults(ctx context.Context, req *cli_servi
7076 return resp , CheckStatus (resp )
7177}
7278
79+ // GetResultSetMetadata is a wrapper around the thrift operation GetResultSetMetadata
80+ // If RecordResults is true, the results will be marshalled to JSON format and written to GetResultSetMetadata<index>.json
7381func (tsc * ThriftServiceClient ) GetResultSetMetadata (ctx context.Context , req * cli_service.TGetResultSetMetadataReq ) (* cli_service.TGetResultSetMetadataResp , error ) {
7482 log := logger .WithContext (driverctx .ConnIdFromContext (ctx ), driverctx .CorrelationIdFromContext (ctx ), SprintGuid (req .OperationHandle .OperationId .GUID ))
7583 defer log .Duration (logger .Track ("GetResultSetMetadata" ))
@@ -85,6 +93,8 @@ func (tsc *ThriftServiceClient) GetResultSetMetadata(ctx context.Context, req *c
8593 return resp , CheckStatus (resp )
8694}
8795
96+ // ExecuteStatement is a wrapper around the thrift operation ExecuteStatement
97+ // If RecordResults is true, the results will be marshalled to JSON format and written to ExecuteStatement<index>.json
8898func (tsc * ThriftServiceClient ) ExecuteStatement (ctx context.Context , req * cli_service.TExecuteStatementReq ) (* cli_service.TExecuteStatementResp , error ) {
8999 msg , start := logger .Track ("ExecuteStatement" )
90100 resp , err := tsc .TCLIServiceClient .ExecuteStatement (context .Background (), req )
@@ -106,6 +116,8 @@ func (tsc *ThriftServiceClient) ExecuteStatement(ctx context.Context, req *cli_s
106116 return resp , CheckStatus (resp )
107117}
108118
119+ // GetOperationStatus is a wrapper around the thrift operation GetOperationStatus
120+ // If RecordResults is true, the results will be marshalled to JSON format and written to GetOperationStatus<index>.json
109121func (tsc * ThriftServiceClient ) GetOperationStatus (ctx context.Context , req * cli_service.TGetOperationStatusReq ) (* cli_service.TGetOperationStatusResp , error ) {
110122 log := logger .WithContext (driverctx .ConnIdFromContext (ctx ), driverctx .CorrelationIdFromContext (ctx ), SprintGuid (req .OperationHandle .OperationId .GUID ))
111123 defer log .Duration (logger .Track ("GetOperationStatus" ))
@@ -121,6 +133,8 @@ func (tsc *ThriftServiceClient) GetOperationStatus(ctx context.Context, req *cli
121133 return resp , CheckStatus (resp )
122134}
123135
136+ // CloseOperation is a wrapper around the thrift operation CloseOperation
137+ // If RecordResults is true, the results will be marshalled to JSON format and written to CloseOperation<index>.json
124138func (tsc * ThriftServiceClient ) CloseOperation (ctx context.Context , req * cli_service.TCloseOperationReq ) (* cli_service.TCloseOperationResp , error ) {
125139 log := logger .WithContext (driverctx .ConnIdFromContext (ctx ), driverctx .CorrelationIdFromContext (ctx ), SprintGuid (req .OperationHandle .OperationId .GUID ))
126140 defer log .Duration (logger .Track ("CloseOperation" ))
@@ -136,6 +150,8 @@ func (tsc *ThriftServiceClient) CloseOperation(ctx context.Context, req *cli_ser
136150 return resp , CheckStatus (resp )
137151}
138152
153+ // CancelOperation is a wrapper around the thrift operation CancelOperation
154+ // If RecordResults is true, the results will be marshalled to JSON format and written to CancelOperation<index>.json
139155func (tsc * ThriftServiceClient ) CancelOperation (ctx context.Context , req * cli_service.TCancelOperationReq ) (* cli_service.TCancelOperationResp , error ) {
140156 log := logger .WithContext (driverctx .ConnIdFromContext (ctx ), driverctx .CorrelationIdFromContext (ctx ), SprintGuid (req .OperationHandle .OperationId .GUID ))
141157 defer log .Duration (logger .Track ("CancelOperation" ))
@@ -157,9 +173,8 @@ func (tsc *ThriftServiceClient) CancelOperation(ctx context.Context, req *cli_se
157173// log.Debug().Msg(c.transport.response.Header.Get("x-thriftserver-error-message"))
158174// log.Debug().Msg(c.transport.response.Header.Get("x-databricks-reason-phrase"))
159175
160- // This is a wrapper of the http transport so we can have access to response code and headers
176+ // InitThriftClient is a wrapper of the http transport, so we can have access to response code and headers.
161177// It is important to know the code and headers to know if we need to retry or not
162-
163178func InitThriftClient (cfg * config.Config ) (* ThriftServiceClient , error ) {
164179 endpoint := cfg .ToEndpointURL ()
165180 tcfg := & thrift.TConfiguration {
@@ -224,11 +239,13 @@ func InitThriftClient(cfg *config.Config) (*ThriftServiceClient, error) {
224239 return tsClient , nil
225240}
226241
227- // ThriftResponse respresents thrift rpc response
242+ // ThriftResponse represents the thrift rpc response
228243type ThriftResponse interface {
229244 GetStatus () * cli_service.TStatus
230245}
231246
247+ // CheckStatus checks the status code after a thrift operation.
248+ // Returns nil if the operation is successful or still executing, otherwise returns an error.
232249func CheckStatus (resp interface {}) error {
233250 rpcresp , ok := resp .(ThriftResponse )
234251 if ok {
@@ -247,6 +264,7 @@ func CheckStatus(resp interface{}) error {
247264 return errors .New ("thrift: invalid response" )
248265}
249266
267+ // SprintGuid is a convenience function to format a byte array into GUID.
250268func SprintGuid (bts []byte ) string {
251269 if len (bts ) == 16 {
252270 return fmt .Sprintf ("%x-%x-%x-%x-%x" , bts [0 :4 ], bts [4 :6 ], bts [6 :8 ], bts [8 :10 ], bts [10 :16 ])
0 commit comments