Skip to content

Commit 7c2ae24

Browse files
iancardosozupiancardosozup
authored
database: feature - add RawQuery function on write_interface.go (#160)
Signed-off-by: Ian Cardoso <[email protected]>
1 parent f6a11a5 commit 7c2ae24

File tree

5 files changed

+135
-13
lines changed

5 files changed

+135
-13
lines changed

‎pkg/entities/vulnerability/vulnerability.go‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -40,17 +40,18 @@ type Vulnerability struct {
4040
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"`
4141
Severity severities.Severity `json:"severity" gorm:"Column:severity" example:"CRITICAL" enums:"CRITICAL, HIGH, MEDIUM, LOW, INFO"`
4242
Type vulnerability.Type `json:"type" gorm:"Column:type" example:"Vulnerability" enums:"Vulnerability, Risk Accepted, False Positive, Corrected"`
43-
CWEs []string `json:"-" gorm:"-" example:"[\"https://cwe.mitre.org/data/definitions/000.html\"]"`
44-
CVEs []string `json:"-" gorm:"-" example:"[\"https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-0000-00000\"]"`
45-
Mitigation string `json:"-" gorm:"-" example:"Use a secret manager or environment variable"`
46-
Reference string `json:"-" gorm:"-" example:"https://example.com"`
47-
SafeExample string `json:"-" gorm:"-" example:"setPassword(env.get(\"HORUSEC_PASSWORD\"))"`
48-
UnsafeExample string `json:"-" gorm:"-" example:"setPassword(\"s@f3P@a$$w0rd\")"`
49-
CommitAuthor string `json:"commitAuthor" gorm:"Column:commit_author" example:"horusec"`
50-
CommitEmail string `json:"commitEmail" gorm:"Column:commit_email" example:"[email protected]"`
51-
CommitHash string `json:"commitHash" gorm:"Column:commit_hash" example:"a21fa164c00a15f3e91f5ee6659cb6a793b39a8d"`
52-
CommitMessage string `json:"commitMessage" gorm:"Column:commit_message" example:"Initial commit"`
53-
CommitDate string `json:"commitDate" gorm:"Column:commit_date" example:"2021-12-30"`
43+
44+
CWEs []string `json:"-" gorm:"-" example:"[\"https://cwe.mitre.org/data/definitions/000.html\"]"`
45+
CVEs []string `json:"-" gorm:"-" example:"[\"https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-0000-00000\"]"`
46+
Mitigation string `json:"-" gorm:"-" example:"Use a secret manager or environment variable"`
47+
Reference string `json:"-" gorm:"-" example:"https://example.com"`
48+
SafeExample string `json:"-" gorm:"-" example:"setPassword(env.get(\"HORUSEC_PASSWORD\"))"`
49+
UnsafeExample string `json:"-" gorm:"-" example:"setPassword(\"s@f3P@a$$w0rd\")"`
50+
CommitAuthor string `json:"commitAuthor" gorm:"Column:commit_author" example:"horusec"`
51+
CommitEmail string `json:"commitEmail" gorm:"Column:commit_email" example:"[email protected]"`
52+
CommitHash string `json:"commitHash" gorm:"Column:commit_hash" example:"a21fa164c00a15f3e91f5ee6659cb6a793b39a8d"`
53+
CommitMessage string `json:"commitMessage" gorm:"Column:commit_message" example:"Initial commit"`
54+
CommitDate string `json:"commitDate" gorm:"Column:commit_date" example:"2021-12-30"`
5455

5556
// RuleID is the rule id used to generate Vulnerability.
5657
// This field can bem empty if Vulnerability was not generated from horusec-engine.
@@ -66,8 +67,8 @@ type Vulnerability struct {
6667
// TODO: This will be removed after the release v2.10.0 of the Horusec CLI be released.
6768
DeprecatedHashes []string `json:"deprecatedHashes" gorm:"-" example:""`
6869

69-
SecurityToolVersion string `json:"securityToolVersion"`
70-
SecurityToolInfoURI string `json:"securityToolInfoUri"`
70+
SecurityToolVersion string `json:"securityToolVersion" gorm:"-"`
71+
SecurityToolInfoURI string `json:"securityToolInfoUri" gorm:"-"`
7172
}
7273

7374
func (v *Vulnerability) GetTable() string {

‎pkg/services/database/database.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,3 +243,23 @@ func (d *database) findPreloadWitLimitAndPageQuery(
243243

244244
return d.connectionRead.Table(table).Where(where).Limit(limit).Offset(page * limit)
245245
}
246+
247+
// Deprecated: Exec starts a transaction and try to execute the raw query into database.
248+
// is not recommended using this and the method will not be available after cli v2.10.0
249+
func (d *database) Exec(rawQuery string, values ...interface{}) error {
250+
sqlDB, err := d.connectionWrite.DB()
251+
if err != nil {
252+
return err
253+
}
254+
255+
tx, err := sqlDB.Begin()
256+
if err != nil {
257+
return err
258+
}
259+
260+
if _, err = tx.Exec(rawQuery, values...); err != nil {
261+
return tx.Rollback()
262+
}
263+
264+
return tx.Commit()
265+
}

‎pkg/services/database/database_mock.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,11 @@ func (m *Mock) Delete(_ map[string]interface{}, _ string) response.IResponse {
8282
return args.Get(0).(response.IResponse)
8383
}
8484

85+
func (m *Mock) Exec(_ string, _ ...interface{}) error {
86+
args := m.MethodCalled("Exec")
87+
return mockUtils.ReturnNilOrError(args, 0)
88+
}
89+
8590
func (m *Mock) FindPreload(entityPointer interface{}, _ map[string]interface{}, _ map[string][]interface{}, _ string) response.IResponse {
8691
args := m.MethodCalled("FindPreload")
8792
return m.reflectValues(entityPointer, args.Get(0).(response.IResponse))

‎pkg/services/database/database_test.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -821,3 +821,98 @@ func TestFindPreloadWitLimitAndPage(t *testing.T) {
821821
assert.Equal(t, nil, response.GetData())
822822
})
823823
}
824+
825+
func TestBatch(t *testing.T) {
826+
t.Run("should success batch query", func(t *testing.T) {
827+
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
828+
assert.NoError(t, err)
829+
query := "UPDATE table a SET a.value = ? where a.id = ? ;"
830+
mock.ExpectBegin()
831+
mock.ExpectExec(query).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
832+
mock.ExpectCommit()
833+
834+
database := &database{
835+
config: config.NewDatabaseConfig(),
836+
connectionRead: getMockedConnection(db),
837+
connectionWrite: getMockedConnection(db),
838+
}
839+
840+
err = database.Exec(query, "1", "2")
841+
842+
assert.NoError(t, err)
843+
})
844+
t.Run("should error when exec fails", func(t *testing.T) {
845+
db, mock, err := sqlmock.New()
846+
assert.NoError(t, err)
847+
query := "UPDATE table a SET a.value = ? where a.id = ? ;"
848+
values := []string{"1", "2"}
849+
mock.ExpectBegin()
850+
mock.ExpectExec(query).WithArgs().WillReturnError(errors.New("some error"))
851+
mock.ExpectCommit()
852+
853+
database := &database{
854+
config: config.NewDatabaseConfig(),
855+
connectionRead: getMockedConnection(db),
856+
connectionWrite: getMockedConnection(db),
857+
}
858+
859+
err = database.Exec(query, values)
860+
861+
assert.Error(t, err)
862+
})
863+
t.Run("should error when begin fails", func(t *testing.T) {
864+
db, mock, err := sqlmock.New()
865+
assert.NoError(t, err)
866+
query := "UPDATE table a SET a.value = ? where a.id = ? ;"
867+
values := []string{"1", "2"}
868+
mock.ExpectBegin().WillReturnError(errors.New("some error"))
869+
mock.ExpectExec(query).WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
870+
mock.ExpectCommit()
871+
872+
database := &database{
873+
config: config.NewDatabaseConfig(),
874+
connectionRead: getMockedConnection(db),
875+
connectionWrite: getMockedConnection(db),
876+
}
877+
878+
err = database.Exec(query, values)
879+
880+
assert.Error(t, err)
881+
})
882+
t.Run("should error when commit fails", func(t *testing.T) {
883+
db, mock, err := sqlmock.New()
884+
assert.NoError(t, err)
885+
query := "UPDATE table a SET a.value = ? where a.id = ? ;"
886+
values := []string{"1", "2"}
887+
mock.ExpectBegin()
888+
mock.ExpectExec(query).WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
889+
mock.ExpectCommit().WillReturnError(errors.New("some error"))
890+
891+
database := &database{
892+
config: config.NewDatabaseConfig(),
893+
connectionRead: getMockedConnection(db),
894+
connectionWrite: getMockedConnection(db),
895+
}
896+
897+
err = database.Exec(query, values)
898+
899+
assert.Error(t, err)
900+
})
901+
t.Run("should error when get DB fails", func(t *testing.T) {
902+
db := &gorm.DB{
903+
Config: &gorm.Config{},
904+
}
905+
906+
query := "UPDATE table a SET a.value = ? where a.id = ? ;"
907+
values := []string{"1", "2"}
908+
909+
database := &database{
910+
config: config.NewDatabaseConfig(),
911+
connectionWrite: db,
912+
}
913+
914+
err := database.Exec(query, values)
915+
916+
assert.Error(t, err)
917+
})
918+
}

‎pkg/services/database/write_interface.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,5 @@ type IDatabaseWrite interface {
2525
CreateOrUpdate(entityPointer interface{}, where map[string]interface{}, table string) response.IResponse
2626
Update(entityPointer interface{}, where map[string]interface{}, table string) response.IResponse
2727
Delete(where map[string]interface{}, table string) response.IResponse
28+
Exec(rawQuery string, values ...interface{}) error
2829
}

0 commit comments

Comments
 (0)