From 7c2ae245b337c95db7ea310522c046ad2394495e Mon Sep 17 00:00:00 2001 From: iancardosozup <86669442+iancardosozup@users.noreply.github.com> Date: Thu, 24 Feb 2022 08:42:13 -0300 Subject: [PATCH] database: feature - add RawQuery function on write_interface.go (#160) Signed-off-by: Ian Cardoso --- pkg/entities/vulnerability/vulnerability.go | 27 +++--- pkg/services/database/database.go | 20 +++++ pkg/services/database/database_mock.go | 5 ++ pkg/services/database/database_test.go | 95 +++++++++++++++++++++ pkg/services/database/write_interface.go | 1 + 5 files changed, 135 insertions(+), 13 deletions(-) diff --git a/pkg/entities/vulnerability/vulnerability.go b/pkg/entities/vulnerability/vulnerability.go index 7fd1ba5..b047de0 100644 --- a/pkg/entities/vulnerability/vulnerability.go +++ b/pkg/entities/vulnerability/vulnerability.go @@ -40,17 +40,18 @@ type Vulnerability struct { Language languages.Language `json:"language" gorm:"Column:language" example:"Leaks" enums:"Go,C#,Dart,Ruby,Python,Java,Kotlin,Javascript,Typescript,Leaks,HCL,C,PHP,HTML,Generic,YAML,Elixir,Shell,Nginx"` Severity severities.Severity `json:"severity" gorm:"Column:severity" example:"CRITICAL" enums:"CRITICAL, HIGH, MEDIUM, LOW, INFO"` Type vulnerability.Type `json:"type" gorm:"Column:type" example:"Vulnerability" enums:"Vulnerability, Risk Accepted, False Positive, Corrected"` - CWEs []string `json:"-" gorm:"-" example:"[\"https://cwe.mitre.org/data/definitions/000.html\"]"` - CVEs []string `json:"-" gorm:"-" example:"[\"https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-0000-00000\"]"` - Mitigation string `json:"-" gorm:"-" example:"Use a secret manager or environment variable"` - Reference string `json:"-" gorm:"-" example:"https://example.com"` - SafeExample string `json:"-" gorm:"-" example:"setPassword(env.get(\"HORUSEC_PASSWORD\"))"` - UnsafeExample string `json:"-" gorm:"-" example:"setPassword(\"s@f3P@a$$w0rd\")"` - CommitAuthor string `json:"commitAuthor" gorm:"Column:commit_author" example:"horusec"` - CommitEmail string `json:"commitEmail" gorm:"Column:commit_email" example:"horusec@zup.com.br"` - CommitHash string `json:"commitHash" gorm:"Column:commit_hash" example:"a21fa164c00a15f3e91f5ee6659cb6a793b39a8d"` - CommitMessage string `json:"commitMessage" gorm:"Column:commit_message" example:"Initial commit"` - CommitDate string `json:"commitDate" gorm:"Column:commit_date" example:"2021-12-30"` + + CWEs []string `json:"-" gorm:"-" example:"[\"https://cwe.mitre.org/data/definitions/000.html\"]"` + CVEs []string `json:"-" gorm:"-" example:"[\"https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-0000-00000\"]"` + Mitigation string `json:"-" gorm:"-" example:"Use a secret manager or environment variable"` + Reference string `json:"-" gorm:"-" example:"https://example.com"` + SafeExample string `json:"-" gorm:"-" example:"setPassword(env.get(\"HORUSEC_PASSWORD\"))"` + UnsafeExample string `json:"-" gorm:"-" example:"setPassword(\"s@f3P@a$$w0rd\")"` + CommitAuthor string `json:"commitAuthor" gorm:"Column:commit_author" example:"horusec"` + CommitEmail string `json:"commitEmail" gorm:"Column:commit_email" example:"horusec@zup.com.br"` + CommitHash string `json:"commitHash" gorm:"Column:commit_hash" example:"a21fa164c00a15f3e91f5ee6659cb6a793b39a8d"` + CommitMessage string `json:"commitMessage" gorm:"Column:commit_message" example:"Initial commit"` + CommitDate string `json:"commitDate" gorm:"Column:commit_date" example:"2021-12-30"` // RuleID is the rule id used to generate Vulnerability. // This field can bem empty if Vulnerability was not generated from horusec-engine. @@ -66,8 +67,8 @@ type Vulnerability struct { // TODO: This will be removed after the release v2.10.0 of the Horusec CLI be released. DeprecatedHashes []string `json:"deprecatedHashes" gorm:"-" example:""` - SecurityToolVersion string `json:"securityToolVersion"` - SecurityToolInfoURI string `json:"securityToolInfoUri"` + SecurityToolVersion string `json:"securityToolVersion" gorm:"-"` + SecurityToolInfoURI string `json:"securityToolInfoUri" gorm:"-"` } func (v *Vulnerability) GetTable() string { diff --git a/pkg/services/database/database.go b/pkg/services/database/database.go index 3107fb2..6a6d6fd 100644 --- a/pkg/services/database/database.go +++ b/pkg/services/database/database.go @@ -243,3 +243,23 @@ func (d *database) findPreloadWitLimitAndPageQuery( return d.connectionRead.Table(table).Where(where).Limit(limit).Offset(page * limit) } + +// Deprecated: Exec starts a transaction and try to execute the raw query into database. +// is not recommended using this and the method will not be available after cli v2.10.0 +func (d *database) Exec(rawQuery string, values ...interface{}) error { + sqlDB, err := d.connectionWrite.DB() + if err != nil { + return err + } + + tx, err := sqlDB.Begin() + if err != nil { + return err + } + + if _, err = tx.Exec(rawQuery, values...); err != nil { + return tx.Rollback() + } + + return tx.Commit() +} diff --git a/pkg/services/database/database_mock.go b/pkg/services/database/database_mock.go index f00281b..722a27b 100644 --- a/pkg/services/database/database_mock.go +++ b/pkg/services/database/database_mock.go @@ -82,6 +82,11 @@ func (m *Mock) Delete(_ map[string]interface{}, _ string) response.IResponse { return args.Get(0).(response.IResponse) } +func (m *Mock) Exec(_ string, _ ...interface{}) error { + args := m.MethodCalled("Exec") + return mockUtils.ReturnNilOrError(args, 0) +} + func (m *Mock) FindPreload(entityPointer interface{}, _ map[string]interface{}, _ map[string][]interface{}, _ string) response.IResponse { args := m.MethodCalled("FindPreload") return m.reflectValues(entityPointer, args.Get(0).(response.IResponse)) diff --git a/pkg/services/database/database_test.go b/pkg/services/database/database_test.go index d633fba..ae887ad 100644 --- a/pkg/services/database/database_test.go +++ b/pkg/services/database/database_test.go @@ -821,3 +821,98 @@ func TestFindPreloadWitLimitAndPage(t *testing.T) { assert.Equal(t, nil, response.GetData()) }) } + +func TestBatch(t *testing.T) { + t.Run("should success batch query", func(t *testing.T) { + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual)) + assert.NoError(t, err) + query := "UPDATE table a SET a.value = ? where a.id = ? ;" + mock.ExpectBegin() + mock.ExpectExec(query).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectCommit() + + database := &database{ + config: config.NewDatabaseConfig(), + connectionRead: getMockedConnection(db), + connectionWrite: getMockedConnection(db), + } + + err = database.Exec(query, "1", "2") + + assert.NoError(t, err) + }) + t.Run("should error when exec fails", func(t *testing.T) { + db, mock, err := sqlmock.New() + assert.NoError(t, err) + query := "UPDATE table a SET a.value = ? where a.id = ? ;" + values := []string{"1", "2"} + mock.ExpectBegin() + mock.ExpectExec(query).WithArgs().WillReturnError(errors.New("some error")) + mock.ExpectCommit() + + database := &database{ + config: config.NewDatabaseConfig(), + connectionRead: getMockedConnection(db), + connectionWrite: getMockedConnection(db), + } + + err = database.Exec(query, values) + + assert.Error(t, err) + }) + t.Run("should error when begin fails", func(t *testing.T) { + db, mock, err := sqlmock.New() + assert.NoError(t, err) + query := "UPDATE table a SET a.value = ? where a.id = ? ;" + values := []string{"1", "2"} + mock.ExpectBegin().WillReturnError(errors.New("some error")) + mock.ExpectExec(query).WithArgs().WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectCommit() + + database := &database{ + config: config.NewDatabaseConfig(), + connectionRead: getMockedConnection(db), + connectionWrite: getMockedConnection(db), + } + + err = database.Exec(query, values) + + assert.Error(t, err) + }) + t.Run("should error when commit fails", func(t *testing.T) { + db, mock, err := sqlmock.New() + assert.NoError(t, err) + query := "UPDATE table a SET a.value = ? where a.id = ? ;" + values := []string{"1", "2"} + mock.ExpectBegin() + mock.ExpectExec(query).WithArgs().WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectCommit().WillReturnError(errors.New("some error")) + + database := &database{ + config: config.NewDatabaseConfig(), + connectionRead: getMockedConnection(db), + connectionWrite: getMockedConnection(db), + } + + err = database.Exec(query, values) + + assert.Error(t, err) + }) + t.Run("should error when get DB fails", func(t *testing.T) { + db := &gorm.DB{ + Config: &gorm.Config{}, + } + + query := "UPDATE table a SET a.value = ? where a.id = ? ;" + values := []string{"1", "2"} + + database := &database{ + config: config.NewDatabaseConfig(), + connectionWrite: db, + } + + err := database.Exec(query, values) + + assert.Error(t, err) + }) +} diff --git a/pkg/services/database/write_interface.go b/pkg/services/database/write_interface.go index 474c0d3..5931425 100644 --- a/pkg/services/database/write_interface.go +++ b/pkg/services/database/write_interface.go @@ -25,4 +25,5 @@ type IDatabaseWrite interface { CreateOrUpdate(entityPointer interface{}, where map[string]interface{}, table string) response.IResponse Update(entityPointer interface{}, where map[string]interface{}, table string) response.IResponse Delete(where map[string]interface{}, table string) response.IResponse + Exec(rawQuery string, values ...interface{}) error }