Increasing code coverage in connection.go

Signed-off-by: Matthew Kim <[email protected]>

Author: mattdeekay

connection.go

@@ -18,16 +18,12 @@ type conn struct {
 	session *cli_service.TOpenSessionResp
 }
 
-// The driver does not really implements prepared statements.
-// Statement ExecContext is the same as connection ExecContext
-// Statement QueryContext is the same as connection QueryContext
+// The driver does not really implement prepared statements.
 func (c *conn) Prepare(query string) (driver.Stmt, error) {
 	return &stmt{conn: c, query: query}, nil
 }
 
-// The driver does not really implements prepared statements.
-// Statement ExecContext is the same as connection ExecContext
-// Statement QueryContext is the same as connection QueryContext
+// The driver does not really implement prepared statements.
 func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
 	return &stmt{conn: c, query: query}, nil
 }
@@ -78,6 +74,7 @@ func (c *conn) IsValid() bool {
 // as an INSERT or UPDATE.
 //
 // ExecContext honors the context timeout and return when it is canceled.
+// Statement ExecContext is the same as connection ExecContext
 func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
 	if len(args) > 0 {
 		return nil, fmt.Errorf("databricks: query parameters are not supported")
@@ -96,6 +93,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 // SELECT.
 //
 // QueryContext honors the context timeout and return when it is canceled.
+// Statement QueryContext is the same as connection QueryContext
 func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
 	if len(args) > 0 {
 		return nil, fmt.Errorf("databricks: query parameters are not supported")

connection_test.go

@@ -874,7 +874,7 @@ func TestConn_QueryContext(t *testing.T) {
 			client:  testClient,
 			cfg:     config.WithDefaults(),
 		}
-		res, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{
+		res, err := testConn.QueryContext(context.Background(), "select 1", []driver.NamedValue{
 			{Value: 1, Name: "name"},
 		})
 
@@ -909,7 +909,7 @@ func TestConn_QueryContext(t *testing.T) {
 			client:  testClient,
 			cfg:     config.WithDefaults(),
 		}
-		res, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{})
+		res, err := testConn.QueryContext(context.Background(), "select 1", []driver.NamedValue{})
 
 		assert.Error(t, err)
 		assert.Nil(t, res)
@@ -960,6 +960,121 @@ func TestConn_QueryContext(t *testing.T) {
 	})
 }
 
+func TestConn_Ping(t *testing.T) {
+	t.Run("ping returns ErrBadConn when executeStatement fails", func(t *testing.T) {
+		var executeStatementCount int
+		executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) {
+			executeStatementCount++
+			executeStatementResp := &cli_service.TExecuteStatementResp{
+				Status: &cli_service.TStatus{
+					StatusCode: cli_service.TStatusCode_ERROR_STATUS,
+				},
+				OperationHandle: &cli_service.TOperationHandle{
+					OperationId: &cli_service.THandleIdentifier{
+						GUID:   []byte("2"),
+						Secret: []byte("b"),
+					},
+				},
+			}
+			return executeStatementResp, nil
+		}
+
+		testClient := &client.TestClient{
+			FnExecuteStatement: executeStatement,
+		}
+		testConn := &conn{
+			session: getTestSession(),
+			client:  testClient,
+			cfg:     config.WithDefaults(),
+		}
+		err := testConn.Ping(context.Background())
+
+		assert.Error(t, err)
+		assert.Equal(t, driver.ErrBadConn, err)
+		assert.Equal(t, 1, executeStatementCount)
+	})
+
+	t.Run("ping returns nil error when driver can establish connection", func(t *testing.T) {
+		var executeStatementCount int
+		executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) {
+			executeStatementCount++
+			executeStatementResp := &cli_service.TExecuteStatementResp{
+				Status: &cli_service.TStatus{
+					StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+				},
+				OperationHandle: &cli_service.TOperationHandle{
+					OperationId: &cli_service.THandleIdentifier{
+						GUID:   []byte("2"),
+						Secret: []byte("b"),
+					},
+				},
+			}
+			return executeStatementResp, nil
+		}
+
+		getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) {
+			getOperationStatusResp := &cli_service.TGetOperationStatusResp{
+				OperationState:  cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE),
+				NumModifiedRows: thrift.Int64Ptr(10),
+			}
+			return getOperationStatusResp, nil
+		}
+
+		testClient := &client.TestClient{
+			FnExecuteStatement:   executeStatement,
+			FnGetOperationStatus: getOperationStatus,
+		}
+
+		testConn := &conn{
+			session: getTestSession(),
+			client:  testClient,
+			cfg:     config.WithDefaults(),
+		}
+		err := testConn.Ping(context.Background())
+
+		assert.Nil(t, err)
+		assert.Equal(t, 1, executeStatementCount)
+	})
+}
+
+func TestConn_Begin(t *testing.T) {
+	t.Run("Begin not supported", func(t *testing.T) {
+		testConn := &conn{
+			session: getTestSession(),
+			client:  &client.TestClient{},
+			cfg:     config.WithDefaults(),
+		}
+		res, err := testConn.Begin()
+		assert.Nil(t, res)
+		assert.Error(t, err)
+	})
+}
+
+func TestConn_BeginTx(t *testing.T) {
+	t.Run("BeginTx not supported", func(t *testing.T) {
+		testConn := &conn{
+			session: getTestSession(),
+			client:  &client.TestClient{},
+			cfg:     config.WithDefaults(),
+		}
+		res, err := testConn.BeginTx(context.Background(), driver.TxOptions{})
+		assert.Nil(t, res)
+		assert.Error(t, err)
+	})
+}
+
+func TestConn_ResetSession(t *testing.T) {
+	t.Run("ResetSession not currently supported", func(t *testing.T) {
+		testConn := &conn{
+			session: getTestSession(),
+			client:  &client.TestClient{},
+			cfg:     config.WithDefaults(),
+		}
+		res := testConn.ResetSession(context.Background())
+		assert.Nil(t, res)
+	})
+}
+
 func getTestSession() *cli_service.TOpenSessionResp {
 	return &cli_service.TOpenSessionResp{SessionHandle: &cli_service.TSessionHandle{
 		SessionId: &cli_service.THandleIdentifier{

statement.go

@@ -34,6 +34,7 @@ func (s *stmt) Query(args []driver.Value) (driver.Rows, error) {
 // as an INSERT or UPDATE.
 //
 // ExecContext honors the context timeout and return when it is canceled.
+// Statement ExecContext is the same as connection ExecContext
 func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
 	return s.conn.ExecContext(ctx, s.query, args)
 }
@@ -42,6 +43,7 @@ func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (drive
 // SELECT.
 //
 // QueryContext honors the context timeout and return when it is canceled.
+// Statement QueryContext is the same as connection QueryContext
 func (s *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
 	return s.conn.QueryContext(ctx, s.query, args)
 }