diff --git a/.github/workflows/check-db-schema-structs.yaml b/.github/workflows/check-db-schema-structs.yaml index 41dba3b0df..b87d31a790 100644 --- a/.github/workflows/check-db-schema-structs.yaml +++ b/.github/workflows/check-db-schema-structs.yaml @@ -51,3 +51,23 @@ jobs: git diff exit 1 fi + check-sqlite-schema-structs: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version: "1.24.4" + - name: Generate SQLite DB schema structs + run: make gen/gorm/sqlite + - name: Check if there are uncommitted file changes + run: | + clean=$(git status --porcelain) + if [[ -z "$clean" ]]; then + echo "SQLite schema is up to date." + else + echo "Uncommitted file changes detected after generating SQLite schema: $clean" + git diff + exit 1 + fi diff --git a/Makefile b/Makefile index 0f22e0c709..07c718051b 100644 --- a/Makefile +++ b/Makefile @@ -114,6 +114,16 @@ start/postgres: stop/postgres: ./scripts/teardown_postgres_db.sh +# Start the SQLite database (file-based, minimal setup) +.PHONY: start/sqlite +start/sqlite: + ./scripts/start_sqlite_db.sh + +# Stop the SQLite database (cleanup) +.PHONY: stop/sqlite +stop/sqlite: + ./scripts/teardown_sqlite_db.sh + # generate the gorm structs for MySQL .PHONY: gen/gorm/mysql gen/gorm/mysql: bin/golang-migrate start/mysql @@ -129,12 +139,22 @@ gen/gorm/postgres: bin/golang-migrate start/postgres cd gorm-gen && GOWORK=off go run main.go --db-type postgres --dsn 'postgres://postgres:postgres@localhost:5432/model-registry?sslmode=disable' && \ cd $(CURDIR) && ./scripts/remove_gorm_defaults.sh) +# generate the gorm structs for SQLite +.PHONY: gen/gorm/sqlite +gen/gorm/sqlite: bin/golang-migrate start/sqlite + @(trap 'cd $(CURDIR) && $(MAKE) stop/sqlite && rm -f /tmp/gorm-gen-sqlite.db' EXIT; \ + $(GOLANG_MIGRATE) -path './internal/datastore/embedmd/sqlite/migrations' -database 'sqlite:///tmp/gorm-gen-sqlite.db' up && \ + cd gorm-gen && GOWORK=off go run main.go --db-type sqlite --dsn '/tmp/gorm-gen-sqlite.db') + # generate the gorm structs (defaults to MySQL for backward compatibility) # Use GORM_DB_TYPE=postgres to generate for PostgreSQL instead +# Use GORM_DB_TYPE=sqlite to generate for SQLite instead .PHONY: gen/gorm gen/gorm: bin/golang-migrate ifeq ($(GORM_DB_TYPE),postgres) $(MAKE) gen/gorm/postgres +else ifeq ($(GORM_DB_TYPE),sqlite) + $(MAKE) gen/gorm/sqlite else $(MAKE) gen/gorm/mysql endif @@ -185,7 +205,7 @@ bin/yq: GOLANG_MIGRATE ?= ${PROJECT_BIN}/migrate bin/golang-migrate: - GOBIN=$(PROJECT_PATH)/bin ${GO} install -tags 'mysql,postgres' github.com/golang-migrate/migrate/v4/cmd/migrate@v4.18.3 + GOBIN=$(PROJECT_PATH)/bin ${GO} install -tags 'mysql,postgres,sqlite' github.com/golang-migrate/migrate/v4/cmd/migrate@v4.18.3 GENQLIENT ?= ${PROJECT_BIN}/genqlient bin/genqlient: diff --git a/go.mod b/go.mod index a56f2a54a9..1f0bfcfea2 100644 --- a/go.mod +++ b/go.mod @@ -28,6 +28,7 @@ require ( google.golang.org/protobuf v1.36.8 gorm.io/driver/mysql v1.6.0 gorm.io/driver/postgres v1.6.0 + gorm.io/driver/sqlite v1.6.0 gorm.io/gorm v1.30.1 k8s.io/api v0.33.4 k8s.io/apimachinery v0.33.4 @@ -102,6 +103,7 @@ require ( github.com/leodido/go-urn v1.4.0 // indirect github.com/lib/pq v1.10.9 // indirect github.com/mailru/easyjson v0.9.0 // indirect + github.com/mattn/go-sqlite3 v1.14.22 // indirect github.com/moby/docker-image-spec v1.3.1 // indirect github.com/moby/go-archive v0.1.0 // indirect github.com/moby/sys/user v0.4.0 // indirect diff --git a/go.sum b/go.sum index 89fc0223e0..358ce85781 100644 --- a/go.sum +++ b/go.sum @@ -303,6 +303,8 @@ github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8S github.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= github.com/mailru/easyjson v0.9.0 h1:PrnmzHw7262yW8sTBwxi1PdJA3Iw/EKBa8psRf7d9a4= github.com/mailru/easyjson v0.9.0/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU= +github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= +github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI= github.com/mdelapenya/tlscert v0.2.0/go.mod h1:O4njj3ELLnJjGdkN7M/vIVCpZ+Cf0L6muqOG4tLSl8o= github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= @@ -661,6 +663,8 @@ gorm.io/driver/mysql v1.6.0 h1:eNbLmNTpPpTOVZi8MMxCi2aaIm0ZpInbORNXDwyLGvg= gorm.io/driver/mysql v1.6.0/go.mod h1:D/oCC2GWK3M/dqoLxnOlaNKmXz8WNTfcS9y5ovaSqKo= gorm.io/driver/postgres v1.6.0 h1:2dxzU8xJ+ivvqTRph34QX+WrRaJlmfyPqXmoGVjMBa4= gorm.io/driver/postgres v1.6.0/go.mod h1:vUw0mrGgrTK+uPHEhAdV4sfFELrByKVGnaVRkXDhtWo= +gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ= +gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8= gorm.io/gorm v1.30.1 h1:lSHg33jJTBxs2mgJRfRZeLDG+WZaHYCk3Wtfl6Ngzo4= gorm.io/gorm v1.30.1/go.mod h1:8Z33v652h4//uMA76KjeDH8mJXPm1QNCYrMeatR0DOE= gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q= diff --git a/go.work.sum b/go.work.sum index 1f5b91ab7b..56cc663d15 100644 --- a/go.work.sum +++ b/go.work.sum @@ -662,6 +662,8 @@ gopkg.in/go-jose/go-jose.v2 v2.6.3 h1:nt80fvSDlhKWQgSWyHyy5CfmlQr+asih51R8PTWNKK gopkg.in/go-jose/go-jose.v2 v2.6.3/go.mod h1:zzZDPkNNw/c9IE7Z9jr11mBZQhKQTMzoEEIoEdZlFBI= gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc= gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc= +gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ= +gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8= honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc h1:/hemPrYIhOhy8zYrNj+069zDB68us2sMGsfkFJO0iZs= istio.io/api v1.24.2 h1:jYjcN6Iq0RPtQj/3KMFsybxmfqmjGN/dxhL7FGJEdIM= istio.io/api v1.24.2/go.mod h1:MQnRok7RZ20/PE56v0LxmoWH0xVxnCQPNuf9O7PAN1I= diff --git a/gorm-gen/main.go b/gorm-gen/main.go index b5d6f63b05..e9533472d8 100644 --- a/gorm-gen/main.go +++ b/gorm-gen/main.go @@ -4,7 +4,6 @@ import ( "fmt" "log" "os" - "strings" "github.com/spf13/cobra" "gorm.io/driver/mysql" @@ -30,15 +29,11 @@ func genModels(g *gen.Generator, db *gorm.DB, tables []string) (err error) { } } - // Custom ModelOpt to remove default tag for nullable fields + // Custom ModelOpt to clean up primary key tags modelOpt := gen.FieldGORMTag("*", func(tag field.GormTag) field.GormTag { - if vals, ok := tag["default"]; ok { - if len(vals) > 0 { - val := strings.Trim(strings.TrimSpace(vals[0]), `"'`) - if strings.ToUpper(val) == "NULL" || val == "0" || val == "" { - tag.Remove("default") - } - } + // Remove "not null" from primary key fields since it's redundant + if _, hasPrimaryKey := tag["primaryKey"]; hasPrimaryKey { + tag.Remove("not null") } return tag }) @@ -46,8 +41,87 @@ func genModels(g *gen.Generator, db *gorm.DB, tables []string) (err error) { // Execute some data table tasks for _, tableName := range tables { if tableName == "Type" { - // Special handling for Type table to set TypeKind as int32 - g.GenerateModel(tableName, gen.FieldType("type_kind", "int32"), modelOpt) + // Special handling for Type table + if dbType == "sqlite" { + g.GenerateModel(tableName, + gen.FieldType("type_kind", "int32"), + gen.FieldType("id", "int32"), + gen.FieldGORMTag("id", func(tag field.GormTag) field.GormTag { + tag.Set("column", "id") + tag.Set("primaryKey", "") + tag.Set("autoIncrement", "true") + return tag + }), + modelOpt) + } else { + g.GenerateModel(tableName, + gen.FieldType("type_kind", "int32"), + gen.FieldGORMTag("id", func(tag field.GormTag) field.GormTag { + tag.Set("column", "id") + tag.Set("primaryKey", "") + tag.Set("autoIncrement", "true") + return tag + }), + modelOpt) + } + } else if tableName == "schema_migrations" && dbType == "sqlite" { + // Special handling for schema_migrations table in SQLite to match MySQL/PostgreSQL types + g.GenerateModel(tableName, + gen.FieldType("version", "int64"), + gen.FieldType("dirty", "bool"), + gen.FieldGORMTag("version", func(tag field.GormTag) field.GormTag { + tag.Set("column", "version") + tag.Set("primaryKey", "") + return tag + }), + gen.FieldGORMTag("dirty", func(tag field.GormTag) field.GormTag { + tag.Set("column", "dirty") + tag.Set("not null", "") + return tag + }), + modelOpt) + } else if isPropertyTable(tableName) { + // Special handling for property tables to ensure composite primary keys + var opts []gen.ModelOpt + if dbType == "sqlite" && hasIDField(tableName) { + opts = append(opts, gen.FieldType("id", "int32")) + opts = append(opts, gen.FieldGORMTag("id", func(tag field.GormTag) field.GormTag { + tag.Set("column", "id") + tag.Set("primaryKey", "") + tag.Set("autoIncrement", "true") + return tag + })) + } + + // Fix composite primary key fields + opts = append(opts, gen.FieldGORMTag("name", func(tag field.GormTag) field.GormTag { + tag.Set("column", "name") + tag.Set("primaryKey", "") + tag.Remove("not null") // Remove not null since primaryKey implies it + return tag + })) + opts = append(opts, gen.FieldGORMTag("is_custom_property", func(tag field.GormTag) field.GormTag { + tag.Set("column", "is_custom_property") + tag.Set("primaryKey", "") + tag.Remove("not null") // Remove not null since primaryKey implies it + return tag + })) + opts = append(opts, modelOpt) + g.GenerateModel(tableName, opts...) + } else if hasIDField(tableName) { + // Tables with ID fields - ensure autoIncrement is set + var opts []gen.ModelOpt + if dbType == "sqlite" { + opts = append(opts, gen.FieldType("id", "int32")) + } + opts = append(opts, gen.FieldGORMTag("id", func(tag field.GormTag) field.GormTag { + tag.Set("column", "id") + tag.Set("primaryKey", "") + tag.Set("autoIncrement", "true") + return tag + })) + opts = append(opts, modelOpt) + g.GenerateModel(tableName, opts...) } else { g.GenerateModel(tableName, modelOpt) } @@ -55,6 +129,35 @@ func genModels(g *gen.Generator, db *gorm.DB, tables []string) (err error) { return nil } +// hasIDField returns true if the table has an auto-increment ID primary key field +func hasIDField(tableName string) bool { + tablesWithIDField := []string{ + "Artifact", "Association", "Attribution", "Context", + "Event", "Execution", "Type", + } + + for _, table := range tablesWithIDField { + if table == tableName { + return true + } + } + return false +} + +// isPropertyTable returns true if the table is a property table with composite primary keys +func isPropertyTable(tableName string) bool { + propertyTables := []string{ + "ArtifactProperty", "ContextProperty", "ExecutionProperty", "TypeProperty", + } + + for _, table := range propertyTables { + if table == tableName { + return true + } + } + return false +} + // getDialector returns the appropriate GORM dialector based on database type and DSN func getDialector(dbType, dsn string) (gorm.Dialector, error) { switch dbType { diff --git a/internal/datastore/embedmd/service.go b/internal/datastore/embedmd/service.go index 18cef47f4b..d9eaa2dda4 100644 --- a/internal/datastore/embedmd/service.go +++ b/internal/datastore/embedmd/service.go @@ -20,8 +20,8 @@ type EmbedMDConfig struct { } func (c *EmbedMDConfig) Validate() error { - if c.DatabaseType != types.DatabaseTypeMySQL && c.DatabaseType != types.DatabaseTypePostgres { - return fmt.Errorf("unsupported database type: %s. Supported types: %s, %s", c.DatabaseType, types.DatabaseTypeMySQL, types.DatabaseTypePostgres) + if c.DatabaseType != types.DatabaseTypeMySQL && c.DatabaseType != types.DatabaseTypePostgres && c.DatabaseType != types.DatabaseTypeSQLite { + return fmt.Errorf("unsupported database type: %s. Supported types: %s, %s, %s", c.DatabaseType, types.DatabaseTypeMySQL, types.DatabaseTypePostgres, types.DatabaseTypeSQLite) } return nil diff --git a/internal/datastore/embedmd/sqlite/connect.go b/internal/datastore/embedmd/sqlite/connect.go new file mode 100644 index 0000000000..e841d7f3b8 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/connect.go @@ -0,0 +1,99 @@ +package sqlite + +import ( + "fmt" + "sync" + "time" + + "github.com/golang/glog" + _tls "github.com/kubeflow/model-registry/internal/tls" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +const ( + // sqliteMaxRetriesDefault is the maximum number of attempts to retry SQLite connection. + sqliteMaxRetriesDefault = 5 // SQLite is file-based, so fewer retries needed +) + +type SQLiteDBConnector struct { + DSN string + TLSConfig *_tls.TLSConfig // Not used for SQLite but kept for interface consistency + db *gorm.DB + connectMutex sync.Mutex + maxRetries int +} + +func NewSQLiteDBConnector( + dsn string, + tlsConfig *_tls.TLSConfig, +) *SQLiteDBConnector { + return &SQLiteDBConnector{ + DSN: dsn, + TLSConfig: tlsConfig, + maxRetries: sqliteMaxRetriesDefault, + } +} + +func (c *SQLiteDBConnector) WithMaxRetries(maxRetries int) *SQLiteDBConnector { + c.maxRetries = maxRetries + + return c +} + +func (c *SQLiteDBConnector) Connect() (*gorm.DB, error) { + // Use mutex to ensure only one connection attempt at a time + c.connectMutex.Lock() + defer c.connectMutex.Unlock() + + // If we already have a working connection, return it + if c.db != nil { + return c.db, nil + } + + var db *gorm.DB + var err error + + // Log warning if TLS configuration is specified (not supported by SQLite) + if c.TLSConfig != nil && c.needsTLSConfig() { + glog.Warningf("TLS configuration is not supported for SQLite connections, ignoring TLS settings") + } + + for i := range c.maxRetries { + glog.V(2).Infof("Attempting to connect to SQLite database: %q (attempt %d/%d)", c.DSN, i+1, c.maxRetries) + db, err = gorm.Open(sqlite.Open(c.DSN), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + TranslateError: true, + }) + if err == nil { + break + } + + glog.Warningf("Retrying connection to SQLite (attempt %d/%d): %v", i+1, c.maxRetries, err) + + time.Sleep(time.Duration(i+1) * time.Second) + } + + if err != nil { + return nil, fmt.Errorf("failed to connect to SQLite: %w", err) + } + + glog.Info("Successfully connected to SQLite database") + + c.db = db + + return db, nil +} + +func (c *SQLiteDBConnector) DB() *gorm.DB { + return c.db +} + +func (c *SQLiteDBConnector) needsTLSConfig() bool { + if c.TLSConfig == nil { + return false + } + + return c.TLSConfig.CertPath != "" || c.TLSConfig.KeyPath != "" || c.TLSConfig.RootCertPath != "" || c.TLSConfig.CAPath != "" || c.TLSConfig.Cipher != "" || c.TLSConfig.VerifyServerCert +} \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/connect_test.go b/internal/datastore/embedmd/sqlite/connect_test.go new file mode 100644 index 0000000000..a38cdae569 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/connect_test.go @@ -0,0 +1,462 @@ +package sqlite_test + +import ( + "os" + "path/filepath" + "testing" + "time" + + "github.com/kubeflow/model-registry/internal/datastore/embedmd/sqlite" + _tls "github.com/kubeflow/model-registry/internal/tls" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSQLiteDBConnector_Connect_Basic(t *testing.T) { + t.Run("InMemoryDatabase", func(t *testing.T) { + // Test in-memory SQLite database + connector := sqlite.NewSQLiteDBConnector(":memory:", &_tls.TLSConfig{}) + + db, err := connector.Connect() + require.NoError(t, err) + assert.NotNil(t, db) + + // Test that we can perform a simple query + var result int + err = db.Raw("SELECT 1").Scan(&result).Error + require.NoError(t, err) + assert.Equal(t, 1, result) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck + }) + + t.Run("FileBasedDatabase", func(t *testing.T) { + // Create temporary file for database + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "test.db") + + connector := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + + db, err := connector.Connect() + require.NoError(t, err) + assert.NotNil(t, db) + + // Test basic operations + var result int + err = db.Raw("SELECT 1").Scan(&result).Error + require.NoError(t, err) + assert.Equal(t, 1, result) + + // Verify database file was created + _, err = os.Stat(dbPath) + require.NoError(t, err) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck + }) + + t.Run("ExistingConnection", func(t *testing.T) { + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "existing.db") + + connector := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + + // First connection + db1, err := connector.Connect() + require.NoError(t, err) + + // Second call should return the same connection + db2, err := connector.Connect() + require.NoError(t, err) + assert.Equal(t, db1, db2) + + // Clean up + sqlDB, err := db1.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck + }) +} + +func TestSQLiteDBConnector_TLSHandling(t *testing.T) { + t.Run("EmptyTLSConfig", func(t *testing.T) { + connector := sqlite.NewSQLiteDBConnector(":memory:", &_tls.TLSConfig{}) + + db, err := connector.Connect() + require.NoError(t, err) + assert.NotNil(t, db) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck + }) + + t.Run("NilTLSConfig", func(t *testing.T) { + connector := sqlite.NewSQLiteDBConnector(":memory:", nil) + + db, err := connector.Connect() + require.NoError(t, err) + assert.NotNil(t, db) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck + }) + + t.Run("TLSConfigWithCertificates", func(t *testing.T) { + // SQLite doesn't support TLS, but should log warning and still connect + tempDir := t.TempDir() + + // Create dummy certificate files (content doesn't matter for this test) + certPath := filepath.Join(tempDir, "cert.pem") + keyPath := filepath.Join(tempDir, "key.pem") + caPath := filepath.Join(tempDir, "ca.pem") + + err := os.WriteFile(certPath, []byte("dummy cert"), 0600) + require.NoError(t, err) + err = os.WriteFile(keyPath, []byte("dummy key"), 0600) + require.NoError(t, err) + err = os.WriteFile(caPath, []byte("dummy ca"), 0600) + require.NoError(t, err) + + tlsConfig := &_tls.TLSConfig{ + CertPath: certPath, + KeyPath: keyPath, + RootCertPath: caPath, + VerifyServerCert: true, + Cipher: "TLS_AES_256_GCM_SHA384", + } + + connector := sqlite.NewSQLiteDBConnector(":memory:", tlsConfig) + + // Should still connect successfully despite TLS config + db, err := connector.Connect() + require.NoError(t, err) + assert.NotNil(t, db) + + // Should be able to perform operations + var result int + err = db.Raw("SELECT 1").Scan(&result).Error + require.NoError(t, err) + assert.Equal(t, 1, result) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck + }) +} + +func TestSQLiteDBConnector_DatabaseOperations(t *testing.T) { + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "operations.db") + + connector := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + + db, err := connector.Connect() + require.NoError(t, err) + defer func() { + sqlDB, _ := db.DB() + sqlDB.Close() //nolint:errcheck + }() + + t.Run("CreateTable", func(t *testing.T) { + err := db.Exec(`CREATE TABLE test_table ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL, + value REAL, + is_active INTEGER, + data BLOB + )`).Error + require.NoError(t, err) + }) + + t.Run("InsertData", func(t *testing.T) { + testData := []byte{0x48, 0x65, 0x6c, 0x6c, 0x6f} // "Hello" in bytes + + err := db.Exec(`INSERT INTO test_table (name, value, is_active, data) + VALUES (?, ?, ?, ?)`, "Test Entry", 3.14159, 1, testData).Error + require.NoError(t, err) + + err = db.Exec(`INSERT INTO test_table (name, value, is_active, data) + VALUES (?, ?, ?, ?)`, "Another Entry", 2.71828, 0, nil).Error + require.NoError(t, err) + }) + + t.Run("QueryData", func(t *testing.T) { + var entries []struct { + ID int64 `gorm:"column:id"` + Name string `gorm:"column:name"` + Value float64 `gorm:"column:value"` + IsActive int `gorm:"column:is_active"` + Data []byte `gorm:"column:data"` + } + + err := db.Raw("SELECT id, name, value, is_active, data FROM test_table ORDER BY id").Scan(&entries).Error + require.NoError(t, err) + require.Len(t, entries, 2) + + assert.Equal(t, "Test Entry", entries[0].Name) + assert.InDelta(t, 3.14159, entries[0].Value, 0.00001) + assert.Equal(t, 1, entries[0].IsActive) + assert.Equal(t, []byte{0x48, 0x65, 0x6c, 0x6c, 0x6f}, entries[0].Data) + + assert.Equal(t, "Another Entry", entries[1].Name) + assert.InDelta(t, 2.71828, entries[1].Value, 0.00001) + assert.Equal(t, 0, entries[1].IsActive) + assert.Nil(t, entries[1].Data) + }) + + t.Run("UpdateData", func(t *testing.T) { + err := db.Exec("UPDATE test_table SET is_active = ? WHERE name = ?", 1, "Another Entry").Error + require.NoError(t, err) + + var isActive int + err = db.Raw("SELECT is_active FROM test_table WHERE name = ?", "Another Entry").Scan(&isActive).Error + require.NoError(t, err) + assert.Equal(t, 1, isActive) + }) + + t.Run("DeleteData", func(t *testing.T) { + err := db.Exec("DELETE FROM test_table WHERE name = ?", "Test Entry").Error + require.NoError(t, err) + + var count int64 + err = db.Raw("SELECT COUNT(*) FROM test_table").Scan(&count).Error + require.NoError(t, err) + assert.Equal(t, int64(1), count) + }) +} + +func TestSQLiteDBConnector_Transactions(t *testing.T) { + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "transactions.db") + + connector := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + + db, err := connector.Connect() + require.NoError(t, err) + defer func() { + sqlDB, _ := db.DB() + sqlDB.Close() //nolint:errcheck + }() + + // Create test table + err = db.Exec(`CREATE TABLE transaction_test ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value TEXT + )`).Error + require.NoError(t, err) + + t.Run("CommitTransaction", func(t *testing.T) { + tx := db.Begin() + + err := tx.Exec("INSERT INTO transaction_test (value) VALUES (?)", "committed_value").Error + require.NoError(t, err) + + err = tx.Commit().Error + require.NoError(t, err) + + // Verify data was committed + var value string + err = db.Raw("SELECT value FROM transaction_test WHERE value = ?", "committed_value").Scan(&value).Error + require.NoError(t, err) + assert.Equal(t, "committed_value", value) + }) + + t.Run("RollbackTransaction", func(t *testing.T) { + tx := db.Begin() + + err := tx.Exec("INSERT INTO transaction_test (value) VALUES (?)", "rollback_value").Error + require.NoError(t, err) + + err = tx.Rollback().Error + require.NoError(t, err) + + // Verify data was not committed + var count int64 + err = db.Raw("SELECT COUNT(*) FROM transaction_test WHERE value = ?", "rollback_value").Scan(&count).Error + require.NoError(t, err) + assert.Equal(t, int64(0), count) + }) +} + +func TestSQLiteDBConnector_Connect_ErrorCases(t *testing.T) { + t.Run("InvalidPath", func(t *testing.T) { + // Try to create database in non-existent directory without permission + invalidPath := "/root/nonexistent/directory/test.db" + connector := sqlite.NewSQLiteDBConnector(invalidPath, &_tls.TLSConfig{}) + + db, err := connector.Connect() + // This might or might not fail depending on the system + // SQLite will try to create the database file, but may fail due to permissions + if err != nil { + assert.Error(t, err) + assert.Nil(t, db) + } else { + // If it succeeds, clean up + sqlDB, _ := db.DB() + sqlDB.Close() //nolint:errcheck + } + }) + + t.Run("WithMaxRetries", func(t *testing.T) { + connector := sqlite.NewSQLiteDBConnector(":memory:", &_tls.TLSConfig{}) + connector = connector.WithMaxRetries(1) + + db, err := connector.Connect() + require.NoError(t, err) + assert.NotNil(t, db) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck + }) +} + +func TestSQLiteDBConnector_ConcurrentConnections(t *testing.T) { + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "concurrent.db") + + // Test multiple connectors to the same database file + connector1 := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + connector2 := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + + db1, err := connector1.Connect() + require.NoError(t, err) + + db2, err := connector2.Connect() + require.NoError(t, err) + + // Create table with first connection + err = db1.Exec(`CREATE TABLE IF NOT EXISTS concurrent_test ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + message TEXT + )`).Error + require.NoError(t, err) + + // Insert data with first connection + err = db1.Exec("INSERT INTO concurrent_test (message) VALUES (?)", "from_db1").Error + require.NoError(t, err) + + // Read data with second connection + var message string + err = db2.Raw("SELECT message FROM concurrent_test WHERE message = ?", "from_db1").Scan(&message).Error + require.NoError(t, err) + assert.Equal(t, "from_db1", message) + + // Insert data with second connection + err = db2.Exec("INSERT INTO concurrent_test (message) VALUES (?)", "from_db2").Error + require.NoError(t, err) + + // Read data with first connection + var messages []string + err = db1.Raw("SELECT message FROM concurrent_test ORDER BY id").Scan(&messages).Error + require.NoError(t, err) + assert.Equal(t, []string{"from_db1", "from_db2"}, messages) + + // Clean up + sqlDB1, _ := db1.DB() + sqlDB1.Close() //nolint:errcheck + sqlDB2, _ := db2.DB() + sqlDB2.Close() //nolint:errcheck +} + +func TestSQLiteDBConnector_DB(t *testing.T) { + connector := sqlite.NewSQLiteDBConnector(":memory:", &_tls.TLSConfig{}) + + // Initially should return nil + assert.Nil(t, connector.DB()) + + // After connecting should return the database + db, err := connector.Connect() + require.NoError(t, err) + + retrievedDB := connector.DB() + assert.Equal(t, db, retrievedDB) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck +} + +func TestSQLiteDBConnector_ConnectionRetries(t *testing.T) { + // Test connection retries with a path that might initially fail + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "retry_test.db") + + connector := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + connector = connector.WithMaxRetries(3) + + // This should succeed (SQLite is quite forgiving) + db, err := connector.Connect() + require.NoError(t, err) + assert.NotNil(t, db) + + // Verify database works + var result int + err = db.Raw("SELECT 1").Scan(&result).Error + require.NoError(t, err) + assert.Equal(t, 1, result) + + // Clean up + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.Close() //nolint:errcheck +} + +func TestSQLiteDBConnector_WALMode(t *testing.T) { + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "wal_test.db") + + connector := sqlite.NewSQLiteDBConnector(dbPath, &_tls.TLSConfig{}) + + db, err := connector.Connect() + require.NoError(t, err) + defer func() { + sqlDB, _ := db.DB() + sqlDB.Close() //nolint:errcheck + }() + + // Enable WAL mode for better concurrency + err = db.Exec("PRAGMA journal_mode=WAL").Error + require.NoError(t, err) + + // Verify WAL mode is enabled + var journalMode string + err = db.Raw("PRAGMA journal_mode").Scan(&journalMode).Error + require.NoError(t, err) + assert.Equal(t, "wal", journalMode) + + // Test basic operations in WAL mode + err = db.Exec(`CREATE TABLE wal_test ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + data TEXT + )`).Error + require.NoError(t, err) + + err = db.Exec("INSERT INTO wal_test (data) VALUES (?)", "wal_data").Error + require.NoError(t, err) + + var data string + err = db.Raw("SELECT data FROM wal_test LIMIT 1").Scan(&data).Error + require.NoError(t, err) + assert.Equal(t, "wal_data", data) + + // Verify WAL file exists + walPath := dbPath + "-wal" + time.Sleep(100 * time.Millisecond) // Give SQLite time to create WAL file + _, err = os.Stat(walPath) + // WAL file might not exist immediately or might be cleaned up, so we don't require it + if err == nil { + t.Logf("WAL file created: %s", walPath) + } +} \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrate.go b/internal/datastore/embedmd/sqlite/migrate.go new file mode 100644 index 0000000000..f150119ac9 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrate.go @@ -0,0 +1,86 @@ +package sqlite + +import ( + "embed" + "fmt" + + "github.com/golang-migrate/migrate/v4" + "github.com/golang-migrate/migrate/v4/database/sqlite3" + "github.com/golang-migrate/migrate/v4/source/iofs" + "gorm.io/gorm" +) + +//go:embed migrations/*.sql +var migrations embed.FS + +const ( + MigrationDir = "migrations" +) + +type SQLiteMigrator struct { + migrator *migrate.Migrate +} + +func NewSQLiteMigrator(db *gorm.DB) (*SQLiteMigrator, error) { + sqlDB, err := db.DB() + if err != nil { + return nil, err + } + + driver, err := sqlite3.WithInstance(sqlDB, &sqlite3.Config{}) + if err != nil { + return nil, err + } + + // Create a new source instance from the embedded files + source, err := iofs.New(migrations, MigrationDir) + if err != nil { + return nil, err + } + + m, err := migrate.NewWithInstance( + "iofs", + source, + "sqlite3", + driver, + ) + if err != nil { + return nil, err + } + + return &SQLiteMigrator{ + migrator: m, + }, nil +} + +func (m *SQLiteMigrator) Migrate() error { + if err := m.Up(nil); err != nil && err != migrate.ErrNoChange { + return err + } + + return nil +} + +func (m *SQLiteMigrator) Up(steps *int) error { + if steps == nil { + return m.migrator.Up() + } + + if *steps < 0 { + return fmt.Errorf("steps cannot be negative") + } + + return m.migrator.Steps(*steps) +} + +func (m *SQLiteMigrator) Down(steps *int) error { + if steps == nil { + return m.migrator.Down() + } + + if *steps > 0 { + return fmt.Errorf("steps cannot be positive") + } + + return m.migrator.Steps(*steps) +} \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrate_test.go b/internal/datastore/embedmd/sqlite/migrate_test.go new file mode 100644 index 0000000000..943a5f37e6 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrate_test.go @@ -0,0 +1,270 @@ +package sqlite_test + +import ( + "path/filepath" + "testing" + + "github.com/kubeflow/model-registry/internal/datastore/embedmd/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + sqlitedriver "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +// Type represents the Type table structure +type Type struct { + ID int64 `gorm:"primaryKey"` + Name string `gorm:"column:name"` + Version string `gorm:"column:version"` + ExternalID string `gorm:"column:external_id"` + Description string `gorm:"column:description"` +} + +func (Type) TableName() string { + return "Type" +} + +// TypeProperty represents the TypeProperty table structure +type TypeProperty struct { + ID int64 `gorm:"primaryKey"` + TypeID int64 `gorm:"column:type_id"` + Name string `gorm:"column:name"` + DataType string `gorm:"column:data_type"` + Description string `gorm:"column:description"` +} + +func (TypeProperty) TableName() string { + return "TypeProperty" +} + +// setupTestDB creates a temporary SQLite database for testing +func setupTestDB(t *testing.T) *gorm.DB { + // Create temporary file for SQLite database + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "test.db") + + // Create GORM database connection + db, err := gorm.Open(sqlitedriver.Open(dbPath), &gorm.Config{ + TranslateError: true, + }) + require.NoError(t, err) + + // Ensure cleanup happens properly + t.Cleanup(func() { + sqlDB, err := db.DB() + if err == nil { + _ = sqlDB.Close() + } + }) + + return db +} + +func TestMigrations(t *testing.T) { + db := setupTestDB(t) + + // Create migrator + migrator, err := sqlite.NewSQLiteMigrator(db) + require.NoError(t, err) + + // Run migrations + err = migrator.Migrate() + require.NoError(t, err) + + // Verify MLMDEnv table exists and has expected data + var schemaVersion int + err = db.Raw("SELECT schema_version FROM MLMDEnv LIMIT 1").Scan(&schemaVersion).Error + require.NoError(t, err) + assert.Equal(t, 10, schemaVersion) + + // Verify Type table has expected entries + var count int64 + err = db.Model(&Type{}).Count(&count).Error + require.NoError(t, err) + assert.Greater(t, count, int64(0)) + + // Verify TypeProperty table has expected entries + err = db.Model(&TypeProperty{}).Count(&count).Error + require.NoError(t, err) + assert.Greater(t, count, int64(0)) + + // Verify at least one type exists by checking if any record exists + var typeExists bool + err = db.Raw("SELECT EXISTS(SELECT 1 FROM Type LIMIT 1)").Scan(&typeExists).Error + require.NoError(t, err) + assert.True(t, typeExists) +} + +func TestDownMigrations(t *testing.T) { + db := setupTestDB(t) + + migrator, err := sqlite.NewSQLiteMigrator(db) + require.NoError(t, err) + + // Run migrations first + err = migrator.Migrate() + require.NoError(t, err) + + // Verify tables exist before down migration + var tableCount int64 + err = db.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name != 'schema_migrations'").Scan(&tableCount).Error + require.NoError(t, err) + assert.Greater(t, tableCount, int64(0)) + + // Run down migrations + err = migrator.Down(nil) + require.NoError(t, err) + + // Verify most tables are dropped (should be significantly fewer than before) + var afterTableCount int64 + err = db.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name != 'schema_migrations'").Scan(&afterTableCount).Error + require.NoError(t, err) + // Should have fewer tables than before (complete rollback may not remove all tables) + assert.LessOrEqual(t, afterTableCount, tableCount/2) +} + +func TestStepMigrations(t *testing.T) { + db := setupTestDB(t) + + migrator, err := sqlite.NewSQLiteMigrator(db) + require.NoError(t, err) + + // Test migrating up in steps + steps := 5 + err = migrator.Up(&steps) + require.NoError(t, err) + + // Verify some tables exist but not all + var tableCount int64 + err = db.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name != 'schema_migrations'").Scan(&tableCount).Error + require.NoError(t, err) + assert.Greater(t, tableCount, int64(0)) + assert.Less(t, tableCount, int64(15)) // Should be less than total tables + + // Test migrating down in steps + downSteps := -2 + err = migrator.Down(&downSteps) + require.NoError(t, err) + + // Verify some tables were removed + var newTableCount int64 + err = db.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name != 'schema_migrations'").Scan(&newTableCount).Error + require.NoError(t, err) + assert.Less(t, newTableCount, tableCount) +} + +func TestMigrationValidation(t *testing.T) { + db := setupTestDB(t) + + migrator, err := sqlite.NewSQLiteMigrator(db) + require.NoError(t, err) + + // Test invalid step parameters + t.Run("InvalidUpSteps", func(t *testing.T) { + negativeSteps := -5 + err := migrator.Up(&negativeSteps) + assert.Error(t, err) + assert.Contains(t, err.Error(), "steps cannot be negative") + }) + + t.Run("InvalidDownSteps", func(t *testing.T) { + positiveSteps := 5 + err := migrator.Down(&positiveSteps) + assert.Error(t, err) + assert.Contains(t, err.Error(), "steps cannot be positive") + }) +} + +func TestSQLiteSpecificFeatures(t *testing.T) { + db := setupTestDB(t) + + migrator, err := sqlite.NewSQLiteMigrator(db) + require.NoError(t, err) + + // Run migrations + err = migrator.Migrate() + require.NoError(t, err) + + t.Run("TestSQLiteIntegerPrimaryKey", func(t *testing.T) { + // Verify that our AUTOINCREMENT primary keys work correctly + // Test by inserting and checking auto-increment behavior + err = db.Exec(`INSERT INTO Artifact (type_id, uri) VALUES (1, 'test://uri')`).Error + require.NoError(t, err) + + var maxID int64 + err = db.Raw("SELECT id FROM Artifact ORDER BY id DESC LIMIT 1").Scan(&maxID).Error + require.NoError(t, err) + assert.Greater(t, maxID, int64(0)) + }) + + t.Run("TestSQLiteBooleanHandling", func(t *testing.T) { + // Test boolean values (stored as INTEGER in SQLite) + // Insert a test record with boolean values + err = db.Exec(`INSERT INTO ArtifactProperty + (artifact_id, name, is_custom_property, bool_value) + VALUES (1, 'test_bool', 1, 0)`).Error + require.NoError(t, err) + + // Query back the boolean value + var boolValue bool + err = db.Raw("SELECT bool_value FROM ArtifactProperty WHERE name = 'test_bool'").Scan(&boolValue).Error + require.NoError(t, err) + assert.Equal(t, false, boolValue) // SQLite boolean handled correctly by GORM + }) + + t.Run("TestSQLiteTextHandling", func(t *testing.T) { + // Test TEXT field handling (SQLite's dynamic typing) + err = db.Exec(`INSERT INTO ArtifactProperty + (artifact_id, name, is_custom_property, string_value) + VALUES (2, 'test_string', 0, 'This is a long text string that would be MEDIUMTEXT in MySQL')`).Error + require.NoError(t, err) + + var stringValue string + err = db.Raw("SELECT string_value FROM ArtifactProperty WHERE name = 'test_string'").Scan(&stringValue).Error + require.NoError(t, err) + assert.Contains(t, stringValue, "long text string") + }) + + t.Run("TestSQLiteBlobHandling", func(t *testing.T) { + // Test BLOB field handling + testData := []byte{0x01, 0x02, 0x03, 0x04, 0xFF} + err = db.Exec(`INSERT INTO ArtifactProperty + (artifact_id, name, is_custom_property, byte_value) + VALUES (3, 'test_blob', 0, ?)`, testData).Error + require.NoError(t, err) + + var blobValue []byte + err = db.Raw("SELECT byte_value FROM ArtifactProperty WHERE name = 'test_blob'").Row().Scan(&blobValue) + require.NoError(t, err) + assert.Equal(t, testData, blobValue) + }) +} + +func TestConcurrentAccess(t *testing.T) { + // Create temporary file for SQLite database + tempDir := t.TempDir() + dbPath := filepath.Join(tempDir, "concurrent_test.db") + + // Test that SQLite can handle concurrent connections + db1, err := gorm.Open(sqlitedriver.Open(dbPath), &gorm.Config{ + TranslateError: true, + }) + require.NoError(t, err) + + db2, err := gorm.Open(sqlitedriver.Open(dbPath), &gorm.Config{ + TranslateError: true, + }) + require.NoError(t, err) + + // Run migrations on first connection + migrator1, err := sqlite.NewSQLiteMigrator(db1) + require.NoError(t, err) + err = migrator1.Migrate() + require.NoError(t, err) + + // Verify second connection can read the schema + var tableCount int64 + err = db2.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type='table'").Scan(&tableCount).Error + require.NoError(t, err) + assert.Greater(t, tableCount, int64(0)) +} \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000001_create_artifact_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000001_create_artifact_table.down.sql new file mode 100644 index 0000000000..023e543fc6 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000001_create_artifact_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "Artifact"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000001_create_artifact_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000001_create_artifact_table.up.sql new file mode 100644 index 0000000000..0e9c108d21 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000001_create_artifact_table.up.sql @@ -0,0 +1,18 @@ +-- Create Artifact table for SQLite +CREATE TABLE IF NOT EXISTS "Artifact" ( + id INTEGER PRIMARY KEY, + type_id INTEGER NOT NULL, + uri TEXT, + state INTEGER DEFAULT NULL, + name VARCHAR(255) DEFAULT NULL, + external_id VARCHAR(255) DEFAULT NULL, + create_time_since_epoch BIGINT NOT NULL DEFAULT '0', + last_update_time_since_epoch BIGINT NOT NULL DEFAULT '0', + UNIQUE (external_id), + UNIQUE (type_id, name) +); + +CREATE INDEX idx_artifact_uri ON "Artifact" (uri); +CREATE INDEX idx_artifact_create_time_since_epoch ON "Artifact" (create_time_since_epoch); +CREATE INDEX idx_artifact_last_update_time_since_epoch ON "Artifact" (last_update_time_since_epoch); +CREATE INDEX idx_artifact_external_id ON "Artifact" (external_id); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000002_create_artifactproperty_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000002_create_artifactproperty_table.down.sql new file mode 100644 index 0000000000..04c704844b --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000002_create_artifactproperty_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "ArtifactProperty"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000002_create_artifactproperty_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000002_create_artifactproperty_table.up.sql new file mode 100644 index 0000000000..7e00465a20 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000002_create_artifactproperty_table.up.sql @@ -0,0 +1,17 @@ +-- Create ArtifactProperty table for SQLite +CREATE TABLE IF NOT EXISTS "ArtifactProperty" ( + artifact_id INTEGER NOT NULL, + name TEXT NOT NULL, + is_custom_property BOOLEAN NOT NULL, + int_value INTEGER DEFAULT NULL, + double_value REAL DEFAULT NULL, + string_value TEXT, + byte_value BLOB, + proto_value BLOB, + bool_value BOOLEAN DEFAULT NULL, + PRIMARY KEY (artifact_id, name, is_custom_property) +); + +CREATE INDEX idx_artifact_property_int ON "ArtifactProperty" (name, is_custom_property, int_value); +CREATE INDEX idx_artifact_property_double ON "ArtifactProperty" (name, is_custom_property, double_value); +CREATE INDEX idx_artifact_property_string ON "ArtifactProperty" (name, is_custom_property, string_value); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000003_create_association_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000003_create_association_table.down.sql new file mode 100644 index 0000000000..eb166720b5 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000003_create_association_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "Association"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000003_create_association_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000003_create_association_table.up.sql new file mode 100644 index 0000000000..765372b760 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000003_create_association_table.up.sql @@ -0,0 +1,7 @@ +-- Create Association table for SQLite +CREATE TABLE IF NOT EXISTS "Association" ( + id INTEGER PRIMARY KEY, + context_id INTEGER NOT NULL, + execution_id INTEGER NOT NULL, + UNIQUE (context_id, execution_id) +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000004_create_attribution_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000004_create_attribution_table.down.sql new file mode 100644 index 0000000000..20939a1a80 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000004_create_attribution_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "Attribution"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000004_create_attribution_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000004_create_attribution_table.up.sql new file mode 100644 index 0000000000..d62035ea45 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000004_create_attribution_table.up.sql @@ -0,0 +1,6 @@ +CREATE TABLE IF NOT EXISTS "Attribution" ( + "id" INTEGER PRIMARY KEY, + "context_id" INTEGER NOT NULL, + "artifact_id" INTEGER NOT NULL, + UNIQUE ("context_id","artifact_id") +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000005_create_context_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000005_create_context_table.down.sql new file mode 100644 index 0000000000..a58d5afd3f --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000005_create_context_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "Context"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000005_create_context_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000005_create_context_table.up.sql new file mode 100644 index 0000000000..ccd2935262 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000005_create_context_table.up.sql @@ -0,0 +1,14 @@ +CREATE TABLE IF NOT EXISTS "Context" ( + "id" INTEGER PRIMARY KEY, + "type_id" INTEGER NOT NULL, + "name" TEXT NOT NULL, + "external_id" TEXT DEFAULT NULL, + "create_time_since_epoch" BIGINT NOT NULL DEFAULT 0, + "last_update_time_since_epoch" BIGINT NOT NULL DEFAULT 0, + UNIQUE ("type_id","name"), + UNIQUE ("external_id") +); + +CREATE INDEX "idx_context_create_time_since_epoch" ON "Context" ("create_time_since_epoch"); +CREATE INDEX "idx_context_last_update_time_since_epoch" ON "Context" ("last_update_time_since_epoch"); +CREATE INDEX "idx_context_external_id" ON "Context" ("external_id"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000006_create_contextproperty_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000006_create_contextproperty_table.down.sql new file mode 100644 index 0000000000..86802581a8 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000006_create_contextproperty_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "ContextProperty"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000006_create_contextproperty_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000006_create_contextproperty_table.up.sql new file mode 100644 index 0000000000..96dacdee7c --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000006_create_contextproperty_table.up.sql @@ -0,0 +1,16 @@ +CREATE TABLE IF NOT EXISTS "ContextProperty" ( + "context_id" INTEGER NOT NULL, + "name" TEXT NOT NULL, + "is_custom_property" BOOLEAN NOT NULL, + "int_value" INTEGER DEFAULT NULL, + "double_value" REAL DEFAULT NULL, + "string_value" TEXT, + "byte_value" BLOB, + "proto_value" BLOB, + "bool_value" BOOLEAN DEFAULT NULL, + PRIMARY KEY ("context_id","name","is_custom_property") +); + +CREATE INDEX "idx_context_property_int" ON "ContextProperty" ("name","is_custom_property","int_value"); +CREATE INDEX "idx_context_property_double" ON "ContextProperty" ("name","is_custom_property","double_value"); +CREATE INDEX "idx_context_property_string" ON "ContextProperty" ("name","is_custom_property","string_value"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000007_create_event_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000007_create_event_table.down.sql new file mode 100644 index 0000000000..beb448078c --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000007_create_event_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "Event"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000007_create_event_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000007_create_event_table.up.sql new file mode 100644 index 0000000000..2470e67bcf --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000007_create_event_table.up.sql @@ -0,0 +1,10 @@ +CREATE TABLE IF NOT EXISTS "Event" ( + "id" INTEGER PRIMARY KEY, + "artifact_id" INTEGER NOT NULL, + "execution_id" INTEGER NOT NULL, + "type" INTEGER NOT NULL, + "milliseconds_since_epoch" BIGINT DEFAULT NULL, + UNIQUE ("artifact_id","execution_id","type") +); + +CREATE INDEX "idx_event_execution_id" ON "Event" ("execution_id"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000008_create_eventpath_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000008_create_eventpath_table.down.sql new file mode 100644 index 0000000000..6dd89d17f9 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000008_create_eventpath_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "EventPath"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000008_create_eventpath_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000008_create_eventpath_table.up.sql new file mode 100644 index 0000000000..e7c4248276 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000008_create_eventpath_table.up.sql @@ -0,0 +1,8 @@ +CREATE TABLE IF NOT EXISTS "EventPath" ( + "event_id" INTEGER NOT NULL, + "is_index_step" BOOLEAN NOT NULL, + "step_index" INTEGER DEFAULT NULL, + "step_key" TEXT +); + +CREATE INDEX "idx_eventpath_event_id" ON "EventPath" ("event_id"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000009_create_execution_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000009_create_execution_table.down.sql new file mode 100644 index 0000000000..48afc792ce --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000009_create_execution_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "Execution"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000009_create_execution_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000009_create_execution_table.up.sql new file mode 100644 index 0000000000..15c38d5874 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000009_create_execution_table.up.sql @@ -0,0 +1,15 @@ +CREATE TABLE IF NOT EXISTS "Execution" ( + "id" INTEGER PRIMARY KEY, + "type_id" INTEGER NOT NULL, + "last_known_state" INTEGER DEFAULT NULL, + "name" TEXT DEFAULT NULL, + "external_id" TEXT DEFAULT NULL, + "create_time_since_epoch" BIGINT NOT NULL DEFAULT 0, + "last_update_time_since_epoch" BIGINT NOT NULL DEFAULT 0, + UNIQUE ("external_id"), + UNIQUE ("type_id","name") +); + +CREATE INDEX "idx_execution_create_time_since_epoch" ON "Execution" ("create_time_since_epoch"); +CREATE INDEX "idx_execution_last_update_time_since_epoch" ON "Execution" ("last_update_time_since_epoch"); +CREATE INDEX "idx_execution_external_id" ON "Execution" ("external_id"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000010_create_executionproperty_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000010_create_executionproperty_table.down.sql new file mode 100644 index 0000000000..a42c640176 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000010_create_executionproperty_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "ExecutionProperty"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000010_create_executionproperty_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000010_create_executionproperty_table.up.sql new file mode 100644 index 0000000000..0fc6bb9ecd --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000010_create_executionproperty_table.up.sql @@ -0,0 +1,16 @@ +CREATE TABLE IF NOT EXISTS "ExecutionProperty" ( + "execution_id" INTEGER NOT NULL, + "name" TEXT NOT NULL, + "is_custom_property" BOOLEAN NOT NULL, + "int_value" INTEGER DEFAULT NULL, + "double_value" REAL DEFAULT NULL, + "string_value" TEXT, + "byte_value" BLOB, + "proto_value" BLOB, + "bool_value" BOOLEAN DEFAULT NULL, + PRIMARY KEY ("execution_id","name","is_custom_property") +); + +CREATE INDEX "idx_execution_property_int" ON "ExecutionProperty" ("name","is_custom_property","int_value"); +CREATE INDEX "idx_execution_property_double" ON "ExecutionProperty" ("name","is_custom_property","double_value"); +CREATE INDEX "idx_execution_property_string" ON "ExecutionProperty" ("name","is_custom_property","string_value"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000011_create_mlmdenv_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000011_create_mlmdenv_table.down.sql new file mode 100644 index 0000000000..651cdd9649 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000011_create_mlmdenv_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "MLMDEnv"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000011_create_mlmdenv_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000011_create_mlmdenv_table.up.sql new file mode 100644 index 0000000000..b65ed60679 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000011_create_mlmdenv_table.up.sql @@ -0,0 +1,4 @@ +CREATE TABLE IF NOT EXISTS "MLMDEnv" ( + "schema_version" INTEGER NOT NULL, + PRIMARY KEY ("schema_version") +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000012_create_parentcontext_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000012_create_parentcontext_table.down.sql new file mode 100644 index 0000000000..4e6792d454 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000012_create_parentcontext_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "ParentContext"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000012_create_parentcontext_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000012_create_parentcontext_table.up.sql new file mode 100644 index 0000000000..1e09a4ec23 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000012_create_parentcontext_table.up.sql @@ -0,0 +1,7 @@ +CREATE TABLE IF NOT EXISTS "ParentContext" ( + "context_id" INTEGER NOT NULL, + "parent_context_id" INTEGER NOT NULL, + PRIMARY KEY ("context_id","parent_context_id") +); + +CREATE INDEX "idx_parentcontext_parent_context_id" ON "ParentContext" ("parent_context_id"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000013_create_parenttype_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000013_create_parenttype_table.down.sql new file mode 100644 index 0000000000..0c71d95977 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000013_create_parenttype_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "ParentType"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000013_create_parenttype_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000013_create_parenttype_table.up.sql new file mode 100644 index 0000000000..ca38c61a53 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000013_create_parenttype_table.up.sql @@ -0,0 +1,5 @@ +CREATE TABLE IF NOT EXISTS "ParentType" ( + "type_id" INTEGER NOT NULL, + "parent_type_id" INTEGER NOT NULL, + PRIMARY KEY ("type_id","parent_type_id") +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000014_create_type_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000014_create_type_table.down.sql new file mode 100644 index 0000000000..3088a3984d --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000014_create_type_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "Type"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000014_create_type_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000014_create_type_table.up.sql new file mode 100644 index 0000000000..8ba13f6f06 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000014_create_type_table.up.sql @@ -0,0 +1,14 @@ +CREATE TABLE IF NOT EXISTS "Type" ( + "id" INTEGER PRIMARY KEY, + "name" TEXT NOT NULL, + "version" TEXT DEFAULT NULL, + "type_kind" INTEGER NOT NULL, + "description" TEXT, + "input_type" TEXT, + "output_type" TEXT, + "external_id" TEXT DEFAULT NULL, + UNIQUE ("external_id") +); + +CREATE INDEX "idx_type_name" ON "Type" ("name"); +CREATE INDEX "idx_type_external_id" ON "Type" ("external_id"); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000015_create_typeproperty_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000015_create_typeproperty_table.down.sql new file mode 100644 index 0000000000..03d207df12 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000015_create_typeproperty_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS "TypeProperty"; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000015_create_typeproperty_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000015_create_typeproperty_table.up.sql new file mode 100644 index 0000000000..bce07039e9 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000015_create_typeproperty_table.up.sql @@ -0,0 +1,6 @@ +CREATE TABLE IF NOT EXISTS "TypeProperty" ( + "type_id" INTEGER NOT NULL, + "name" TEXT NOT NULL, + "data_type" INTEGER DEFAULT NULL, + PRIMARY KEY ("type_id","name") +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000016_seed_type_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000016_seed_type_table.down.sql new file mode 100644 index 0000000000..8bf9b9d5bc --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000016_seed_type_table.down.sql @@ -0,0 +1,18 @@ +DELETE FROM "Type" WHERE "name" IN ( + 'mlmd.Dataset', + 'mlmd.Model', + 'mlmd.Metrics', + 'mlmd.Statistics', + 'mlmd.Train', + 'mlmd.Transform', + 'mlmd.Process', + 'mlmd.Evaluate', + 'mlmd.Deploy', + 'kf.RegisteredModel', + 'kf.ModelVersion', + 'kf.DocArtifact', + 'kf.ModelArtifact', + 'kf.ServingEnvironment', + 'kf.InferenceService', + 'kf.ServeModel' +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000016_seed_type_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000016_seed_type_table.up.sql new file mode 100644 index 0000000000..53b2e70a39 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000016_seed_type_table.up.sql @@ -0,0 +1,23 @@ +INSERT INTO "Type" ("name", "version", "type_kind", "description", "input_type", "output_type", "external_id") +SELECT t.* FROM ( + SELECT 'mlmd.Dataset' as name, NULL as version, 1 as type_kind, NULL as description, NULL as input_type, NULL as output_type, NULL as external_id + UNION ALL SELECT 'mlmd.Model', NULL, 1, NULL, NULL, NULL, NULL + UNION ALL SELECT 'mlmd.Metrics', NULL, 1, NULL, NULL, NULL, NULL + UNION ALL SELECT 'mlmd.Statistics', NULL, 1, NULL, NULL, NULL, NULL + UNION ALL SELECT 'mlmd.Train', NULL, 0, NULL, NULL, NULL, NULL + UNION ALL SELECT 'mlmd.Transform', NULL, 0, NULL, NULL, NULL, NULL + UNION ALL SELECT 'mlmd.Process', NULL, 0, NULL, NULL, NULL, NULL + UNION ALL SELECT 'mlmd.Evaluate', NULL, 0, NULL, NULL, NULL, NULL + UNION ALL SELECT 'mlmd.Deploy', NULL, 0, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.RegisteredModel', NULL, 2, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.ModelVersion', NULL, 2, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.DocArtifact', NULL, 1, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.ModelArtifact', NULL, 1, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.ServingEnvironment', NULL, 2, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.InferenceService', NULL, 2, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.ServeModel', NULL, 0, NULL, NULL, NULL, NULL +) t +WHERE NOT EXISTS ( + SELECT 1 FROM "Type" + WHERE "name" = t.name AND "version" IS NULL +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000017_seed_typeproperty_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000017_seed_typeproperty_table.down.sql new file mode 100644 index 0000000000..6deed65644 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000017_seed_typeproperty_table.down.sql @@ -0,0 +1,27 @@ +DELETE FROM "TypeProperty" +WHERE ("type_id", "name", "data_type") IN ( + ((SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'owner', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'state', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'author', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'model_name', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'state', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'version', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.DocArtifact'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'model_format_name', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'model_format_version', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'service_account_name', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'storage_key', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'storage_path', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ServingEnvironment'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'desired_state', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'model_version_id', 1), + ((SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'registered_model_id', 1), + ((SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'runtime', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'serving_environment_id', 1), + ((SELECT id FROM "Type" WHERE name = 'kf.ServeModel'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ServeModel'), 'model_version_id', 1) +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000017_seed_typeproperty_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000017_seed_typeproperty_table.up.sql new file mode 100644 index 0000000000..0955535bfb --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000017_seed_typeproperty_table.up.sql @@ -0,0 +1,25 @@ +INSERT OR IGNORE INTO "TypeProperty" ("type_id", "name", "data_type") +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'owner', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'state', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'author', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'model_name', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'state', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelVersion'), 'version', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.DocArtifact'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'model_format_name', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'model_format_version', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'service_account_name', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'storage_key', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ModelArtifact'), 'storage_path', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ServingEnvironment'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'desired_state', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'model_version_id', 1 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'registered_model_id', 1 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'runtime', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.InferenceService'), 'serving_environment_id', 1 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ServeModel'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ServeModel'), 'model_version_id', 1; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000018_seed_mlmdenv_table.down.sql b/internal/datastore/embedmd/sqlite/migrations/000018_seed_mlmdenv_table.down.sql new file mode 100644 index 0000000000..0145eda212 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000018_seed_mlmdenv_table.down.sql @@ -0,0 +1 @@ +DELETE FROM "MLMDEnv" WHERE "schema_version" = 10; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000018_seed_mlmdenv_table.up.sql b/internal/datastore/embedmd/sqlite/migrations/000018_seed_mlmdenv_table.up.sql new file mode 100644 index 0000000000..06ac114245 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000018_seed_mlmdenv_table.up.sql @@ -0,0 +1,2 @@ +INSERT OR IGNORE INTO "MLMDEnv" ("schema_version") VALUES +(10); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000019_add_base_model_fields_to_registered_model.down.sql b/internal/datastore/embedmd/sqlite/migrations/000019_add_base_model_fields_to_registered_model.down.sql new file mode 100644 index 0000000000..f887f87525 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000019_add_base_model_fields_to_registered_model.down.sql @@ -0,0 +1,13 @@ +DELETE FROM "TypeProperty" WHERE type_id=( + SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel' +) AND "name" IN ( + 'language', + 'library_name', + 'license_link', + 'license', + 'logo', + 'maturity', + 'provider', + 'readme', + 'tasks' +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000019_add_base_model_fields_to_registered_model.up.sql b/internal/datastore/embedmd/sqlite/migrations/000019_add_base_model_fields_to_registered_model.up.sql new file mode 100644 index 0000000000..fe1525a9fe --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000019_add_base_model_fields_to_registered_model.up.sql @@ -0,0 +1,10 @@ +INSERT OR IGNORE INTO "TypeProperty" ("type_id", "name", "data_type") +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'language', 4 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'library_name', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'license_link', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'license', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'logo', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'maturity', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'provider', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'readme', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.RegisteredModel'), 'tasks', 4; \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000020_seed_experiment_types.down.sql b/internal/datastore/embedmd/sqlite/migrations/000020_seed_experiment_types.down.sql new file mode 100644 index 0000000000..f09d4b8a8d --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000020_seed_experiment_types.down.sql @@ -0,0 +1,8 @@ +DELETE FROM "Type" WHERE "name" IN ( + 'kf.MetricHistory', + 'kf.Experiment', + 'kf.ExperimentRun', + 'kf.DataSet', + 'kf.Metric', + 'kf.Parameter' +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000020_seed_experiment_types.up.sql b/internal/datastore/embedmd/sqlite/migrations/000020_seed_experiment_types.up.sql new file mode 100644 index 0000000000..f3dcc374a9 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000020_seed_experiment_types.up.sql @@ -0,0 +1,13 @@ +INSERT INTO "Type" ("name", "version", "type_kind", "description", "input_type", "output_type", "external_id") +SELECT t.* FROM ( + SELECT 'kf.MetricHistory' as name, NULL as version, 1 as type_kind, NULL as description, NULL as input_type, NULL as output_type, NULL as external_id + UNION ALL SELECT 'kf.Experiment', NULL, 2, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.ExperimentRun', NULL, 2, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.DataSet', NULL, 1, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.Metric', NULL, 1, NULL, NULL, NULL, NULL + UNION ALL SELECT 'kf.Parameter', NULL, 1, NULL, NULL, NULL, NULL +) t +WHERE NOT EXISTS ( + SELECT 1 FROM "Type" + WHERE "name" = t.name AND "version" IS NULL +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000021_add_experiment_type_properties.down.sql b/internal/datastore/embedmd/sqlite/migrations/000021_add_experiment_type_properties.down.sql new file mode 100644 index 0000000000..df36653193 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000021_add_experiment_type_properties.down.sql @@ -0,0 +1,30 @@ +DELETE FROM "TypeProperty" +WHERE ("type_id", "name", "data_type") IN ( + ((SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'value', 5), + ((SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'timestamp', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'step', 1), + ((SELECT id FROM "Type" WHERE name = 'kf.Parameter'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.Parameter'), 'value', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.Parameter'), 'parameter_type', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'value', 5), + ((SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'timestamp', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'step', 1), + ((SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'digest', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'source_type', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'source', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'schema', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'profile', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.Experiment'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.Experiment'), 'owner', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.Experiment'), 'state', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'description', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'owner', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'state', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'status', 3), + ((SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'start_time_since_epoch', 1), + ((SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'end_time_since_epoch', 1), + ((SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'experiment_id', 1) +); \ No newline at end of file diff --git a/internal/datastore/embedmd/sqlite/migrations/000021_add_experiment_type_properties.up.sql b/internal/datastore/embedmd/sqlite/migrations/000021_add_experiment_type_properties.up.sql new file mode 100644 index 0000000000..05ee7c3f86 --- /dev/null +++ b/internal/datastore/embedmd/sqlite/migrations/000021_add_experiment_type_properties.up.sql @@ -0,0 +1,28 @@ +INSERT OR IGNORE INTO "TypeProperty" ("type_id", "name", "data_type") +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'value', 5 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'timestamp', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Metric'), 'step', 1 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Parameter'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Parameter'), 'value', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Parameter'), 'parameter_type', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'value', 5 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'timestamp', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.MetricHistory'), 'step', 1 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'digest', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'source_type', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'source', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'schema', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.DataSet'), 'profile', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Experiment'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Experiment'), 'owner', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.Experiment'), 'state', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'description', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'owner', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'state', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'status', 3 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'start_time_since_epoch', 1 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'end_time_since_epoch', 1 UNION ALL +SELECT (SELECT id FROM "Type" WHERE name = 'kf.ExperimentRun'), 'experiment_id', 1; \ No newline at end of file diff --git a/internal/db/connector.go b/internal/db/connector.go index 2d088ffd9f..c05e57f092 100644 --- a/internal/db/connector.go +++ b/internal/db/connector.go @@ -6,6 +6,7 @@ import ( "github.com/kubeflow/model-registry/internal/datastore/embedmd/mysql" "github.com/kubeflow/model-registry/internal/datastore/embedmd/postgres" + "github.com/kubeflow/model-registry/internal/datastore/embedmd/sqlite" "github.com/kubeflow/model-registry/internal/db/types" "github.com/kubeflow/model-registry/internal/tls" "gorm.io/gorm" @@ -34,8 +35,10 @@ func Init(dbType string, dsn string, tlsConfig *tls.TLSConfig) error { _connectorInstance = mysql.NewMySQLDBConnector(dsn, tlsConfig) case "postgres": _connectorInstance = postgres.NewPostgresDBConnector(dsn, tlsConfig) + case "sqlite": + _connectorInstance = sqlite.NewSQLiteDBConnector(dsn, tlsConfig) default: - return fmt.Errorf("unsupported database type: %s. Supported types: %s, %s", dbType, types.DatabaseTypeMySQL, types.DatabaseTypePostgres) + return fmt.Errorf("unsupported database type: %s. Supported types: %s, %s, %s", dbType, types.DatabaseTypeMySQL, types.DatabaseTypePostgres, types.DatabaseTypeSQLite) } return nil diff --git a/internal/db/migrator.go b/internal/db/migrator.go index 90973ddba2..892032efc0 100644 --- a/internal/db/migrator.go +++ b/internal/db/migrator.go @@ -5,6 +5,7 @@ import ( "github.com/kubeflow/model-registry/internal/datastore/embedmd/mysql" "github.com/kubeflow/model-registry/internal/datastore/embedmd/postgres" + "github.com/kubeflow/model-registry/internal/datastore/embedmd/sqlite" "github.com/kubeflow/model-registry/internal/db/types" "gorm.io/gorm" ) @@ -21,7 +22,9 @@ func NewDBMigrator(dbType string, db *gorm.DB) (DBMigrator, error) { return mysql.NewMySQLMigrator(db) case types.DatabaseTypePostgres: return postgres.NewPostgresMigrator(db) + case types.DatabaseTypeSQLite: + return sqlite.NewSQLiteMigrator(db) } - return nil, fmt.Errorf("unsupported database type: %s. Supported types: %s, %s", dbType, types.DatabaseTypeMySQL, types.DatabaseTypePostgres) + return nil, fmt.Errorf("unsupported database type: %s. Supported types: %s, %s, %s", dbType, types.DatabaseTypeMySQL, types.DatabaseTypePostgres, types.DatabaseTypeSQLite) } diff --git a/internal/db/types/types.go b/internal/db/types/types.go index b5057ed039..6cbe67cc51 100644 --- a/internal/db/types/types.go +++ b/internal/db/types/types.go @@ -5,4 +5,6 @@ const ( DatabaseTypeMySQL = "mysql" // DatabaseTypePostgres represents PostgreSQL database type DatabaseTypePostgres = "postgres" + // DatabaseTypeSQLite represents SQLite database type + DatabaseTypeSQLite = "sqlite" ) diff --git a/scripts/start_sqlite_db.sh b/scripts/start_sqlite_db.sh new file mode 100755 index 0000000000..b4f8c22d82 --- /dev/null +++ b/scripts/start_sqlite_db.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash + +# SQLite database setup script +# Since SQLite is file-based, this script mainly ensures the database directory exists +# and provides a consistent interface with other database setup scripts + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(dirname "${SCRIPT_DIR}")" + +# Default SQLite database path (can be overridden by environment variable) +# Use temp directory for CI, project directory for local development +if [ "${CI:-false}" = "true" ]; then + SQLITE_DB_PATH="${SQLITE_DB_PATH:-/tmp/model-registry.db}" +else + SQLITE_DB_PATH="${SQLITE_DB_PATH:-${PROJECT_DIR}/model-registry.db}" +fi + +echo "Starting SQLite database setup..." +echo "Database file: ${SQLITE_DB_PATH}" + +# Create the directory for the database file if it doesn't exist +DB_DIR="$(dirname "${SQLITE_DB_PATH}")" +if [ ! -d "${DB_DIR}" ]; then + echo "Creating database directory: ${DB_DIR}" + mkdir -p "${DB_DIR}" +fi + +# SQLite will create the database file automatically when first accessed +# For consistency with other database scripts, we can touch the file +if [ ! -f "${SQLITE_DB_PATH}" ]; then + echo "Creating empty database file: ${SQLITE_DB_PATH}" + touch "${SQLITE_DB_PATH}" +fi + +echo "SQLite database setup completed successfully." +echo "Database file: ${SQLITE_DB_PATH}" +echo "" +echo "To connect to the database manually, use:" +echo " sqlite3 ${SQLITE_DB_PATH}" +echo "" +echo "To run migrations, use:" +echo " make gen/gorm/sqlite" \ No newline at end of file diff --git a/scripts/teardown_sqlite_db.sh b/scripts/teardown_sqlite_db.sh new file mode 100755 index 0000000000..e0906e9f64 --- /dev/null +++ b/scripts/teardown_sqlite_db.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash + +# SQLite database teardown script +# Removes the SQLite database file and any related temporary files + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(dirname "${SCRIPT_DIR}")" + +# Default SQLite database path (can be overridden by environment variable) +SQLITE_DB_PATH="${SQLITE_DB_PATH:-${PROJECT_DIR}/model-registry.db}" + +echo "Tearing down SQLite database..." +echo "Database file: ${SQLITE_DB_PATH}" + +# Remove the main database file if it exists +if [ -f "${SQLITE_DB_PATH}" ]; then + echo "Removing database file: ${SQLITE_DB_PATH}" + rm -f "${SQLITE_DB_PATH}" +else + echo "Database file does not exist: ${SQLITE_DB_PATH}" +fi + +# Remove any SQLite temporary files (WAL, SHM files) +SQLITE_WAL_PATH="${SQLITE_DB_PATH}-wal" +SQLITE_SHM_PATH="${SQLITE_DB_PATH}-shm" + +if [ -f "${SQLITE_WAL_PATH}" ]; then + echo "Removing WAL file: ${SQLITE_WAL_PATH}" + rm -f "${SQLITE_WAL_PATH}" +fi + +if [ -f "${SQLITE_SHM_PATH}" ]; then + echo "Removing SHM file: ${SQLITE_SHM_PATH}" + rm -f "${SQLITE_SHM_PATH}" +fi + +echo "SQLite database teardown completed successfully." \ No newline at end of file