From 94f087bae3592ca33f716be94355c779e7547b6f Mon Sep 17 00:00:00 2001
From: andream16 <andrea@ocurity.com>
Date: Tue, 26 Nov 2024 15:13:35 +0000
Subject: [PATCH] Automigrate on local SQLite backend.

---
 .../storer/local/sqlite/export_test.go        | 25 -----------------
 .../internal/storer/local/sqlite/sqlite.go    | 27 +++++++++++++++++++
 .../storer/local/sqlite/sqlite_test.go        |  3 ---
 3 files changed, 27 insertions(+), 28 deletions(-)
 delete mode 100644 sdk/component/internal/storer/local/sqlite/export_test.go

diff --git a/sdk/component/internal/storer/local/sqlite/export_test.go b/sdk/component/internal/storer/local/sqlite/export_test.go
deleted file mode 100644
index 02fcc135d..000000000
--- a/sdk/component/internal/storer/local/sqlite/export_test.go
+++ /dev/null
@@ -1,25 +0,0 @@
-package sqlite
-
-import "github.com/go-errors/errors"
-
-// CreateTable is used to create a table in testing settings.
-func (m *manager) CreateTable() error {
-	stmt, err := m.db.Prepare(`
-		CREATE TABLE finding (
-			id INTEGER PRIMARY KEY AUTOINCREMENT UNIQUE,
-			instance_id UUID NOT NULL UNIQUE,
-			findings TEXT NOT NULL,
-			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-			updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-		);
-	`)
-	if err != nil {
-		return errors.Errorf("could not prepare statement for creating table: %w", err)
-	}
-
-	if _, err := stmt.Exec(); err != nil {
-		return errors.Errorf("could not create table: %w", err)
-	}
-
-	return stmt.Close()
-}
diff --git a/sdk/component/internal/storer/local/sqlite/sqlite.go b/sdk/component/internal/storer/local/sqlite/sqlite.go
index 20a2f72ff..a05a1ab95 100644
--- a/sdk/component/internal/storer/local/sqlite/sqlite.go
+++ b/sdk/component/internal/storer/local/sqlite/sqlite.go
@@ -74,6 +74,10 @@ func NewManager(dsn string, opts ...managerOption) (*manager, error) {
 		}
 	}
 
+	if err := mgr.migrate(); err != nil {
+		return nil, errors.Errorf("could not apply migrations: %w", err)
+	}
+
 	return mgr, nil
 }
 
@@ -211,6 +215,29 @@ func (m *manager) Close(ctx context.Context) error {
 	return nil
 }
 
+// TODO: potentially leverage migrations here but this is simple enough for now for local setup.
+// Tracked here https://linear.app/smithy/issue/OCU-274/automigrate-on-sqlite-storage.
+func (m *manager) migrate() error {
+	stmt, err := m.db.Prepare(`
+		CREATE TABLE IF NOT EXISTS finding (
+			id INTEGER PRIMARY KEY AUTOINCREMENT UNIQUE,
+			instance_id UUID NOT NULL UNIQUE,
+			findings TEXT NOT NULL,
+			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+			updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+		);
+	`)
+	if err != nil {
+		return fmt.Errorf("could not prepare statement for creating table: %w", err)
+	}
+
+	if _, err := stmt.Exec(); err != nil {
+		return fmt.Errorf("could not create table: %w", err)
+	}
+
+	return stmt.Close()
+}
+
 func (m *manager) marshalFindings(findings []*ocsf.VulnerabilityFinding) (string, error) {
 	var rawFindings []json.RawMessage
 	for _, finding := range findings {
diff --git a/sdk/component/internal/storer/local/sqlite/sqlite_test.go b/sdk/component/internal/storer/local/sqlite/sqlite_test.go
index bd483276b..2f50dc323 100644
--- a/sdk/component/internal/storer/local/sqlite/sqlite_test.go
+++ b/sdk/component/internal/storer/local/sqlite/sqlite_test.go
@@ -29,8 +29,6 @@ type (
 		component.Reader
 		component.Updater
 		component.Writer
-
-		CreateTable() error
 	}
 
 	ManagerTestSuite struct {
@@ -54,7 +52,6 @@ func (mts *ManagerTestSuite) SetupTest() {
 
 	mts.manager, err = sqlite.NewManager("smithy.db", sqlite.ManagerWithClock(clock))
 	require.NoError(mts.t, err)
-	require.NoError(mts.T(), mts.manager.CreateTable())
 }
 
 func (mts *ManagerTestSuite) TearDownTest() {