database: feature - add RawQuery function on write_interface.go (#160)

Signed-off-by: Ian Cardoso <ian.cardoso@zup.com.br>
ZupIT · Feb 24, 2022 · 7c2ae24 · 7c2ae24
1 parent f6a11a5
commit 7c2ae24
Show file tree

Hide file tree

Showing 5 changed files with 135 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
 }