Skip to content

Commit

Permalink
updates
Browse files Browse the repository at this point in the history
  • Loading branch information
caffix committed Dec 12, 2024
1 parent a458ed8 commit 626b130
Show file tree
Hide file tree
Showing 5 changed files with 42 additions and 17 deletions.
3 changes: 1 addition & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@ module github.com/owasp-amass/asset-db
go 1.23.1

require (
github.com/caffix/queue v0.3.1
github.com/caffix/stringset v0.2.0
github.com/glebarez/sqlite v1.11.0
github.com/neo4j/neo4j-go-driver/v5 v5.27.0
github.com/owasp-amass/open-asset-model v0.12.0
github.com/rubenv/sql-migrate v1.7.0
github.com/stretchr/testify v1.9.0
Expand Down Expand Up @@ -33,7 +33,6 @@ require (
github.com/ncruces/go-strftime v0.1.9 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/stretchr/objx v0.5.2 // indirect
golang.org/x/crypto v0.29.0 // indirect
golang.org/x/exp v0.0.0-20241108190413-2d47ceb2692f // indirect
golang.org/x/sync v0.9.0 // indirect
Expand Down
6 changes: 2 additions & 4 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
github.com/caffix/queue v0.3.1 h1:eg4V7cGomH76/Xq60MdEhSVZvO6E3iOPlacMwZOjiSA=
github.com/caffix/queue v0.3.1/go.mod h1:AnCrUsy3cwDjLJvNv7FkQx/0Q9Hb7n7agRrUP0OiYl4=
github.com/caffix/stringset v0.2.0 h1:kN6xnvL8jzx2YhQNOYr6A6hFzUK+iikt1JtJ2MS2LC8=
github.com/caffix/stringset v0.2.0/go.mod h1:8PZ6GIPpMP5+r5hr790/05w3v9xI+gXRxRzJCZL57lQ=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
Expand Down Expand Up @@ -52,6 +50,8 @@ github.com/microsoft/go-mssqldb v0.17.0 h1:Fto83dMZPnYv1Zwx5vHHxpNraeEaUlQ/hhHLg
github.com/microsoft/go-mssqldb v0.17.0/go.mod h1:OkoNGhGEs8EZqchVTtochlXruEhEOaO4S0d2sB5aeGQ=
github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4=
github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/neo4j/neo4j-go-driver/v5 v5.27.0 h1:YdsIxDjAQbjlP/4Ha9B/gF8Y39UdgdTwCyihSxy8qTw=
github.com/neo4j/neo4j-go-driver/v5 v5.27.0/go.mod h1:Vff8OwT7QpLm7L2yYr85XNWe9Rbqlbeb9asNXJTHO4k=
github.com/owasp-amass/open-asset-model v0.12.0 h1:WBf0P82ONVJErGjIdZ9jkAzXssOT87dFfY6dKM93csc=
github.com/owasp-amass/open-asset-model v0.12.0/go.mod h1:DOX+SiD6PZBroSMnsILAmpf0SHi6TVpqjV4uNfBeg7g=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
Expand All @@ -65,8 +65,6 @@ github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/f
github.com/rubenv/sql-migrate v1.7.0 h1:HtQq1xyTN2ISmQDggnh0c9U3JlP8apWh8YO2jzlXpTI=
github.com/rubenv/sql-migrate v1.7.0/go.mod h1:S4wtDEG1CKn+0ShpTtzWhFpHHI5PvCUtiGI+C+Z2THE=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
Expand Down
37 changes: 31 additions & 6 deletions repository/neo4j/db.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,31 +6,56 @@ package neo4j

import (
"context"
"fmt"
"net/url"
"strings"
"time"

neo4jdb "github.com/neo4j/neo4j-go-driver/v5/neo4j"
"github.com/neo4j/neo4j-go-driver/v5/neo4j/config"
)

const Neo4j string = "neo4j"

// neoRepository is a repository implementation using Neo4j as the underlying DBMS.
type neoRepository struct {
db neo4jdb.DriverWithContext
db neo4jdb.DriverWithContext
dbname string
}

// New creates a new instance of the asset database repository.
func New(dbtype, dsn string) (*neoRepository, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
u, err := url.Parse(dsn)
if err != nil {
return nil, err
}

driver, err := neo4jdb.NewDriverWithContext(dsn, neo4jdb.NoAuth(), func(config *neo4jdb.Config) {
config.MaxConnectionPoolSize = 10
auth := neo4jdb.NoAuth()
var username, password string
if u.User != nil {
username = u.User.Username()
password, _ = u.User.Password()
auth = neo4jdb.BasicAuth(username, password, "")
}
dbname := strings.TrimPrefix(u.Path, "/")

newdsn := fmt.Sprintf("%s://%s", u.Scheme, u.Host)
driver, err := neo4jdb.NewDriverWithContext(newdsn, auth, func(cfg *config.Config) {
cfg.MaxConnectionPoolSize = 20
cfg.MaxConnectionLifetime = time.Hour
cfg.ConnectionLivenessCheckTimeout = 10 * time.Minute
})
if err != nil {
return nil, err
}

return &neoRepository{db: driver}, driver.VerifyConnectivity(ctx)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()

return &neoRepository{
db: driver,
dbname: dbname,
}, driver.VerifyConnectivity(ctx)
}

// Close implements the Repository interface.
Expand Down
3 changes: 3 additions & 0 deletions repository/repository.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"strings"
"time"

"github.com/owasp-amass/asset-db/repository/neo4j"
"github.com/owasp-amass/asset-db/repository/sqlrepo"
"github.com/owasp-amass/asset-db/types"
oam "github.com/owasp-amass/open-asset-model"
Expand Down Expand Up @@ -47,6 +48,8 @@ type Repository interface {
// New creates a new instance of the asset database repository.
func New(dbtype, dsn string) (Repository, error) {
switch strings.ToLower(dbtype) {
case strings.ToLower(neo4j.Neo4j):
return neo4j.New(dbtype, dsn)

Check failure on line 52 in repository/repository.go

View workflow job for this annotation

GitHub Actions / integration

cannot use neo4j.New(dbtype, dsn) (value of type *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository) as Repository value in return statement: *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository does not implement Repository (missing method CreateAsset)

Check failure on line 52 in repository/repository.go

View workflow job for this annotation

GitHub Actions / lint

cannot use neo4j.New(dbtype, dsn) (value of type *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository) as Repository value in return statement: *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository does not implement Repository (missing method CreateAsset) (typecheck)

Check failure on line 52 in repository/repository.go

View workflow job for this annotation

GitHub Actions / lint

cannot use neo4j.New(dbtype, dsn) (value of type *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository) as Repository value in return statement: *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository does not implement Repository (missing method CreateAsset)) (typecheck)

Check failure on line 52 in repository/repository.go

View workflow job for this annotation

GitHub Actions / lint

cannot use neo4j.New(dbtype, dsn) (value of type *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository) as Repository value in return statement: *"github.com/owasp-amass/asset-db/repository/neo4j".neoRepository does not implement Repository (missing method CreateAsset)) (typecheck)
case strings.ToLower(sqlrepo.Postgres):
fallthrough
case strings.ToLower(sqlrepo.SQLite):
Expand Down
10 changes: 5 additions & 5 deletions repository/sqlrepo/db.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,9 @@ func newDatabase(dbtype, dsn string) (*gorm.DB, error) {
case Postgres:
return postgresDatabase(dsn)
case SQLite:
fallthrough
return sqliteDatabase(dsn, 3, 5)
case SQLiteMemory:
return sqliteDatabase(dsn)
return sqliteDatabase(dsn, 50, 100)
}
return nil, errors.New("unknown DB type")
}
Expand All @@ -72,7 +72,7 @@ func postgresDatabase(dsn string) (*gorm.DB, error) {
}

// sqliteDatabase creates a new SQLite database connection using the provided data source name (dsn).
func sqliteDatabase(dsn string) (*gorm.DB, error) {
func sqliteDatabase(dsn string, conns, idles int) (*gorm.DB, error) {
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
if err != nil {
return nil, err
Expand All @@ -83,8 +83,8 @@ func sqliteDatabase(dsn string) (*gorm.DB, error) {
return nil, err
}

sqlDB.SetMaxIdleConns(5)
sqlDB.SetMaxOpenConns(3)
sqlDB.SetMaxOpenConns(conns)
sqlDB.SetMaxIdleConns(idles)
sqlDB.SetConnMaxLifetime(time.Hour)
sqlDB.SetConnMaxIdleTime(10 * time.Minute)
return db, nil
Expand Down

0 comments on commit 626b130

Please sign in to comment.