This is an automated email from the ASF dual-hosted git repository. maxyang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/cloudberry-go-libs.git
commit cf68aaf5471b901dc78706d9fd8ef1f083399c31 Author: Jamie McAtamney <[email protected]> AuthorDate: Wed Aug 30 03:56:41 2023 -0700 dbconn: Add SelectInt and SelectIntSlice Selecting ints can have the same issues as selecting strings (may be NULL, may not exist, needs a destination struct, etc.) so for user convenience this commit adds SelectInt and SelectIntSlice functions that work essentially like SelectString and SelectStringSlice. --- dbconn/dbconn.go | 55 ++++++++++++++++++++++++++++++++++++++++++ dbconn/dbconn_test.go | 67 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/dbconn/dbconn.go b/dbconn/dbconn.go index 64e2c45..e98ceb2 100644 --- a/dbconn/dbconn.go +++ b/dbconn/dbconn.go @@ -464,3 +464,58 @@ func SelectStringSlice(connection *DBConn, query string, whichConn ...int) ([]st } return retval, nil } + +/* + * The below are convenience functions for selecting one or more ints that may + * be NULL or may not exist, with the same functionality (and the same rationale) + * as SelectString and SelectStringSlice; see the comments for those functions, + * above, for more details. + */ +func MustSelectInt(connection *DBConn, query string, whichConn ...int) int { + str, err := SelectInt(connection, query, whichConn...) + gplog.FatalOnError(err) + return str +} + +func SelectInt(connection *DBConn, query string, whichConn ...int) (int, error) { + results, err := SelectIntSlice(connection, query, whichConn...) + if err != nil { + return 0, err + } + if len(results) == 1 { + return results[0], nil + } else if len(results) > 1 { + return 0, errors.Errorf("Too many rows returned from query: got %d rows, expected 1 row", len(results)) + } + return 0, nil +} + +func MustSelectIntSlice(connection *DBConn, query string, whichConn ...int) []int { + str, err := SelectIntSlice(connection, query, whichConn...) + gplog.FatalOnError(err) + return str +} + +func SelectIntSlice(connection *DBConn, query string, whichConn ...int) ([]int, error) { + connNum := connection.ValidateConnNum(whichConn...) + rows, err := connection.Query(query, connNum) + if err != nil { + return []int{}, err + } + if cols, _ := rows.Rows.Columns(); len(cols) > 1 { + return []int{}, errors.Errorf("Too many columns returned from query: got %d columns, expected 1 column", len(cols)) + } + retval := make([]int, 0) + for rows.Rows.Next() { + var result sql.NullInt32 + err = rows.Rows.Scan(&result) + if err != nil { + return []int{}, err + } + retval = append(retval, int(result.Int32)) + } + if rows.Rows.Err() != nil { + return []int{}, rows.Rows.Err() + } + return retval, nil +} diff --git a/dbconn/dbconn_test.go b/dbconn/dbconn_test.go index e5d342e..b58c67b 100644 --- a/dbconn/dbconn_test.go +++ b/dbconn/dbconn_test.go @@ -480,4 +480,71 @@ var _ = Describe("dbconn/dbconn tests", func() { dbconn.MustSelectString(connection, "SELECT foo FROM bar") }) }) + Describe("MustSelectInt", func() { + header := []string{"foo"} + rowOne := []driver.Value{"1"} + rowTwo := []driver.Value{"2"} + headerExtraCol := []string{"foo", "bar"} + rowExtraCol := []driver.Value{"1", "2"} + + It("returns a single int if the query selects a single int", func() { + fakeResult := sqlmock.NewRows(header).AddRow(rowOne...) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + result := dbconn.MustSelectInt(connection, "SELECT foo FROM bar") + Expect(result).To(Equal(1)) + }) + It("returns 0 if the query selects no ints", func() { + fakeResult := sqlmock.NewRows(header) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + result := dbconn.MustSelectInt(connection, "SELECT foo FROM bar") + Expect(result).To(Equal(0)) + }) + It("panics if the query selects multiple rows", func() { + fakeResult := sqlmock.NewRows(header).AddRow(rowOne...).AddRow(rowTwo...) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + defer testhelper.ShouldPanicWithMessage("Too many rows returned from query: got 2 rows, expected 1 row") + dbconn.MustSelectInt(connection, "SELECT foo FROM bar") + }) + It("panics if the query selects multiple columns", func() { + fakeResult := sqlmock.NewRows(headerExtraCol).AddRow(rowExtraCol...) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + defer testhelper.ShouldPanicWithMessage("Too many columns returned from query: got 2 columns, expected 1 column") + dbconn.MustSelectInt(connection, "SELECT foo FROM bar") + }) + }) + Describe("MustSelectIntSlice", func() { + header := []string{"foo"} + rowOne := []driver.Value{"1"} + rowTwo := []driver.Value{"2"} + headerExtraCol := []string{"foo", "bar"} + rowExtraCol := []driver.Value{"1", "2"} + + It("returns a slice containing a single int if the query selects a single int", func() { + fakeResult := sqlmock.NewRows(header).AddRow(rowOne...) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + results := dbconn.MustSelectIntSlice(connection, "SELECT foo FROM bar") + Expect(len(results)).To(Equal(1)) + Expect(results[0]).To(Equal(1)) + }) + It("returns an empty slice if the query selects no ints", func() { + fakeResult := sqlmock.NewRows(header) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + results := dbconn.MustSelectIntSlice(connection, "SELECT foo FROM bar") + Expect(len(results)).To(Equal(0)) + }) + It("returns a slice containing multiple ints if the query selects multiple rows", func() { + fakeResult := sqlmock.NewRows(header).AddRow(rowOne...).AddRow(rowTwo...) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + results := dbconn.MustSelectIntSlice(connection, "SELECT foo FROM bar") + Expect(len(results)).To(Equal(2)) + Expect(results[0]).To(Equal(1)) + Expect(results[1]).To(Equal(2)) + }) + It("panics if the query selects multiple columns", func() { + fakeResult := sqlmock.NewRows(headerExtraCol).AddRow(rowExtraCol...) + mock.ExpectQuery("SELECT (.*)").WillReturnRows(fakeResult) + defer testhelper.ShouldPanicWithMessage("Too many columns returned from query: got 2 columns, expected 1 column") + dbconn.MustSelectInt(connection, "SELECT foo FROM bar") + }) + }) }) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
