From ce1f0eae3149f4a9d169aa8f51d34bb3e5794aa1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Legin?= <45884629+meeehow@users.noreply.github.com> Date: Mon, 4 Sep 2023 15:36:05 +0200 Subject: [PATCH 1/5] Create postgres tables if they don't exist --- exporters/postgres/postgres.go | 97 +++++++++++++++++++++++++++++++++- storage/postgres/postgres.go | 51 +++++++++++++++++- 2 files changed, 146 insertions(+), 2 deletions(-) diff --git a/exporters/postgres/postgres.go b/exporters/postgres/postgres.go index 7bce7e7..50352b5 100644 --- a/exporters/postgres/postgres.go +++ b/exporters/postgres/postgres.go @@ -48,8 +48,84 @@ func (e *Exporter) Name() string { return Name } -// NewExporter creates new Postregre exporter. +// NewExporter creates new Postregre exporter and all the necessary tables, if they don't exist. func NewExporter(sqlDB *sql.DB, uploadPayloads bool) (*Exporter, error) { + // Check if the "samples" table exists. + exists, err := tableExists(sqlDB, "samples") + if err != nil { + return nil, fmt.Errorf("error while checking if samples table exists: %v", err) + } + + if !exists { + sql := `CREATE TABLE samples ( + sha256 VARCHAR(100) PRIMARY KEY, + mimetype text, + file_output text, + size INT + )` + _, err = sqlDB.Exec(sql) + if err != nil { + return nil, fmt.Errorf("error while creating samples table: %v", err) + } + } + + // Check if the "payloads" table exists. + exists, err = tableExists(sqlDB, "payloads") + if err != nil { + return nil, fmt.Errorf("error while checking if payloads table exists: %v", err) + } + + if !exists { + sql := `CREATE TABLE payloads ( + sha256 VARCHAR(100) PRIMARY KEY, + payload bytea + )` + _, err = sqlDB.Exec(sql) + if err != nil { + return nil, fmt.Errorf("error while creating payloads table: %v", err) + } + } + + // Check if the "sources" table exists. + exists, err = tableExists(sqlDB, "sources") + if err != nil { + return nil, fmt.Errorf("error while checking if sources table exists: %v", err) + } + + if !exists { + sql := `CREATE TABLE sources ( + sha256 VARCHAR(100) PRIMARY KEY, + sourceID text[], + sourcePath text, + sourceDescription text, + repoName text, + repoPath text + )` + _, err = sqlDB.Exec(sql) + if err != nil { + return nil, fmt.Errorf("error while creating sources table: %v", err) + } + } + + // Check if the "samples_sources" table exists. + exists, err = tableExists(sqlDB, "samples_sources") + if err != nil { + return nil, fmt.Errorf("error while checking if samples_sources table exists: %v", err) + } + + if !exists { + sql := `CREATE TABLE samples_sources ( + sample_sha256 VARCHAR(100) REFERENCES samples(sha256) NOT NULL, + source_sha256 VARCHAR(100) REFERENCES sources(sha256) NOT NULL, + sample_paths text[], + PRIMARY KEY (sample_sha256, source_sha256) + )` + _, err = sqlDB.Exec(sql) + if err != nil { + return nil, fmt.Errorf("error while creating samples_sources table: %v", err) + } + } + return &Exporter{sqlDB: sqlDB, uploadPayloads: uploadPayloads}, nil } @@ -283,3 +359,22 @@ func fileCmdOutput(filepath string) (string, error) { return strings.TrimSuffix(stdout.String(), "\n"), nil } + +func tableExists(db *sql.DB, tableName string) (bool, error) { + // Query to check if the table exists in PostgreSQL + query := ` + SELECT EXISTS ( + SELECT 1 + FROM information_schema.tables + WHERE table_name = $1 + ) + ` + + var exists bool + err := db.QueryRow(query, tableName).Scan(&exists) + if err != nil { + return false, err + } + + return exists, nil +} diff --git a/storage/postgres/postgres.go b/storage/postgres/postgres.go index 0241ecc..126b4e9 100644 --- a/storage/postgres/postgres.go +++ b/storage/postgres/postgres.go @@ -18,6 +18,7 @@ package postgres import ( "context" "database/sql" + "fmt" "github.com/google/hashr/core/hashr" @@ -30,8 +31,37 @@ type Storage struct { sqlDB *sql.DB } -// NewStorage creates new Storage struct that allows to interact with PostgreSQL instance. +// NewStorage creates new Storage struct that allows to interact with PostgreSQL instance and all the necessary tables, if they don't exist. func NewStorage(sqlDB *sql.DB) (*Storage, error) { + // Check if the "jobs" table exists. + exists, err := tableExists(sqlDB, "jobs") + if err != nil { + return nil, fmt.Errorf("error while checking if jobs table exists: %v", err) + } + + if !exists { + sql := `CREATE TABLE jobs ( + quick_sha256 VARCHAR(100) PRIMARY KEY, + imported_at INT NOT NULL, + id text, + repo text, + repo_path text, + location text, + sha256 VARCHAR(100), + status VARCHAR(50), + error text, + preprocessing_duration INT, + processing_duration INT, + export_duration INT, + files_extracted INT, + files_exported INT + )` + _, err = sqlDB.Exec(sql) + if err != nil { + return nil, fmt.Errorf("error while creating jobs table: %v", err) + } + } + return &Storage{sqlDB: sqlDB}, nil } @@ -98,3 +128,22 @@ func (s *Storage) FetchJobs(ctx context.Context) (map[string]string, error) { return processed, nil } + +func tableExists(db *sql.DB, tableName string) (bool, error) { + // Query to check if the table exists in PostgreSQL + query := ` + SELECT EXISTS ( + SELECT 1 + FROM information_schema.tables + WHERE table_name = $1 + ) + ` + + var exists bool + err := db.QueryRow(query, tableName).Scan(&exists) + if err != nil { + return false, err + } + + return exists, nil +} \ No newline at end of file From 031bd1b7367679dc2aea12da1b7bf22c3f19abe7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Legin?= <45884629+meeehow@users.noreply.github.com> Date: Tue, 5 Sep 2023 10:14:24 +0200 Subject: [PATCH 2/5] Adjust tests for postgres exporter --- exporters/postgres/postgres_test.go | 5 +++++ storage/postgres/postgres.go | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/exporters/postgres/postgres_test.go b/exporters/postgres/postgres_test.go index 608803e..3084b8e 100644 --- a/exporters/postgres/postgres_test.go +++ b/exporters/postgres/postgres_test.go @@ -33,6 +33,11 @@ func TestExport(t *testing.T) { } defer db.Close() + mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("samples").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("payloads").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("samples_sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) + postgresExporter, err := NewExporter(db, false) if err != nil { t.Fatalf("could not create Postgres exporter: %v", err) diff --git a/storage/postgres/postgres.go b/storage/postgres/postgres.go index 126b4e9..4b981bc 100644 --- a/storage/postgres/postgres.go +++ b/storage/postgres/postgres.go @@ -31,7 +31,7 @@ type Storage struct { sqlDB *sql.DB } -// NewStorage creates new Storage struct that allows to interact with PostgreSQL instance and all the necessary tables, if they don't exist. +// NewStorage creates new Storage struct that allows to interact with PostgreSQL instance and all the necessary tables, if they don't exist.git commit -m func NewStorage(sqlDB *sql.DB) (*Storage, error) { // Check if the "jobs" table exists. exists, err := tableExists(sqlDB, "jobs") From 1144d67b7aa0693fff6139185d1c14a4961ec068 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Legin?= <45884629+meeehow@users.noreply.github.com> Date: Tue, 5 Sep 2023 10:27:44 +0200 Subject: [PATCH 3/5] Adjust expected query --- exporters/postgres/postgres_test.go | 8 ++++---- storage/postgres/postgres.go | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/exporters/postgres/postgres_test.go b/exporters/postgres/postgres_test.go index 3084b8e..288cb2d 100644 --- a/exporters/postgres/postgres_test.go +++ b/exporters/postgres/postgres_test.go @@ -33,10 +33,10 @@ func TestExport(t *testing.T) { } defer db.Close() - mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("samples").WillReturnRows(mock.NewRows([]string{"TRUE"})) - mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("payloads").WillReturnRows(mock.NewRows([]string{"TRUE"})) - mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) - mock.ExpectQuery(`SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)`).WithArgs("samples_sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("payloads").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples_sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) postgresExporter, err := NewExporter(db, false) if err != nil { diff --git a/storage/postgres/postgres.go b/storage/postgres/postgres.go index 4b981bc..b2a6854 100644 --- a/storage/postgres/postgres.go +++ b/storage/postgres/postgres.go @@ -31,7 +31,7 @@ type Storage struct { sqlDB *sql.DB } -// NewStorage creates new Storage struct that allows to interact with PostgreSQL instance and all the necessary tables, if they don't exist.git commit -m +// NewStorage creates new Storage struct that allows to interact with PostgreSQL instance and all the necessary tables, if they don't exist. func NewStorage(sqlDB *sql.DB) (*Storage, error) { // Check if the "jobs" table exists. exists, err := tableExists(sqlDB, "jobs") @@ -146,4 +146,4 @@ func tableExists(db *sql.DB, tableName string) (bool, error) { } return exists, nil -} \ No newline at end of file +} From 1f3edd42be622592575e5efed5237ad7634f2763 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Legin?= <45884629+meeehow@users.noreply.github.com> Date: Tue, 5 Sep 2023 10:51:41 +0200 Subject: [PATCH 4/5] Adjust expected query --- exporters/postgres/postgres_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/exporters/postgres/postgres_test.go b/exporters/postgres/postgres_test.go index 288cb2d..de18b99 100644 --- a/exporters/postgres/postgres_test.go +++ b/exporters/postgres/postgres_test.go @@ -33,10 +33,10 @@ func TestExport(t *testing.T) { } defer db.Close() - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples").WillReturnRows(mock.NewRows([]string{"TRUE"})) - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("payloads").WillReturnRows(mock.NewRows([]string{"TRUE"})) - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples_sources").WillReturnRows(mock.NewRows([]string{"TRUE"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples").WillReturnRows(mock.NewRows([]string{"t"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("payloads").WillReturnRows(mock.NewRows([]string{"t"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("sources").WillReturnRows(mock.NewRows([]string{"t"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples_sources").WillReturnRows(mock.NewRows([]string{"t"})) postgresExporter, err := NewExporter(db, false) if err != nil { From 709e2cffbd1290f0e180ea633380cd1f107d71f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Legin?= <45884629+meeehow@users.noreply.github.com> Date: Tue, 5 Sep 2023 12:41:52 +0200 Subject: [PATCH 5/5] Adjust expected query --- exporters/postgres/postgres.go | 5 +++-- exporters/postgres/postgres_test.go | 8 ++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/exporters/postgres/postgres.go b/exporters/postgres/postgres.go index 50352b5..908f1a7 100644 --- a/exporters/postgres/postgres.go +++ b/exporters/postgres/postgres.go @@ -366,11 +366,12 @@ func tableExists(db *sql.DB, tableName string) (bool, error) { SELECT EXISTS ( SELECT 1 FROM information_schema.tables - WHERE table_name = $1 - ) + WHERE table_name=$1 + ); ` var exists bool + err := db.QueryRow(query, tableName).Scan(&exists) if err != nil { return false, err diff --git a/exporters/postgres/postgres_test.go b/exporters/postgres/postgres_test.go index de18b99..f4052ff 100644 --- a/exporters/postgres/postgres_test.go +++ b/exporters/postgres/postgres_test.go @@ -33,10 +33,10 @@ func TestExport(t *testing.T) { } defer db.Close() - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples").WillReturnRows(mock.NewRows([]string{"t"})) - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("payloads").WillReturnRows(mock.NewRows([]string{"t"})) - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("sources").WillReturnRows(mock.NewRows([]string{"t"})) - mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name = $1 )`).WithArgs("samples_sources").WillReturnRows(mock.NewRows([]string{"t"})) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name=$1 );`).WithArgs("samples").WillReturnRows(mock.NewRows([]string{"t"}).AddRow("t")) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name=$1 );`).WithArgs("payloads").WillReturnRows(mock.NewRows([]string{"t"}).AddRow("t")) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name=$1 );`).WithArgs("sources").WillReturnRows(mock.NewRows([]string{"t"}).AddRow("t")) + mock.ExpectQuery(`SELECT EXISTS ( SELECT 1 FROM information_schema.tables WHERE table_name=$1 );`).WithArgs("samples_sources").WillReturnRows(mock.NewRows([]string{"t"}).AddRow("t")) postgresExporter, err := NewExporter(db, false) if err != nil {