Skip to content

Commit 955da59

Browse files
authored
feat: create the embedmd datastore type (#1107)
Signed-off-by: Alessio Pragliola <[email protected]>
1 parent 4d361d0 commit 955da59

File tree

48 files changed

+687
-8
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

48 files changed

+687
-8
lines changed

cmd/proxy.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,5 +123,8 @@ func init() {
123123
proxyCmd.Flags().StringVar(&proxyCfg.Datastore.MLMD.Hostname, "mlmd-hostname", proxyCfg.Datastore.MLMD.Hostname, "MLMD hostname")
124124
proxyCmd.Flags().IntVar(&proxyCfg.Datastore.MLMD.Port, "mlmd-port", proxyCfg.Datastore.MLMD.Port, "MLMD port")
125125

126+
proxyCmd.Flags().StringVar(&proxyCfg.Datastore.EmbedMD.DatabaseType, "embedmd-database-type", "mysql", "EmbedMD database type")
127+
proxyCmd.Flags().StringVar(&proxyCfg.Datastore.EmbedMD.DatabaseDSN, "embedmd-database-dsn", "", "EmbedMD database DSN")
128+
126129
proxyCmd.Flags().StringVar(&proxyCfg.Datastore.Type, "datastore-type", proxyCfg.Datastore.Type, "Datastore type")
127130
}

go.mod

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ require (
66
github.com/go-chi/chi/v5 v5.2.1
77
github.com/go-chi/cors v1.2.1
88
github.com/go-logr/logr v1.4.2
9+
github.com/golang-migrate/migrate/v4 v4.18.3
910
github.com/golang/glog v1.2.5
1011
github.com/kserve/kserve v0.15.0
1112
github.com/onsi/ginkgo v1.16.5
@@ -16,9 +17,12 @@ require (
1617
github.com/spf13/viper v1.20.1
1718
github.com/stretchr/testify v1.10.0
1819
github.com/testcontainers/testcontainers-go v0.37.0
20+
github.com/testcontainers/testcontainers-go/modules/mysql v0.37.0
1921
go.uber.org/zap v1.27.0
2022
google.golang.org/grpc v1.72.1
2123
google.golang.org/protobuf v1.36.6
24+
gorm.io/driver/mysql v1.5.7
25+
gorm.io/gorm v1.25.7
2226
k8s.io/api v0.32.3
2327
k8s.io/apimachinery v0.32.3
2428
k8s.io/client-go v0.32.3
@@ -63,6 +67,7 @@ require (
6367
github.com/go-openapi/swag v0.23.1 // indirect
6468
github.com/go-playground/locales v0.14.1 // indirect
6569
github.com/go-playground/universal-translator v0.18.1 // indirect
70+
github.com/go-sql-driver/mysql v1.7.1 // indirect
6671
github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
6772
github.com/go-viper/mapstructure/v2 v2.2.1 // indirect
6873
github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect
@@ -79,6 +84,10 @@ require (
7984
github.com/googleapis/gax-go/v2 v2.14.1 // indirect
8085
github.com/googleapis/google-cloud-go-testing v0.0.0-20210719221736-1c9a4c676720 // indirect
8186
github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 // indirect
87+
github.com/hashicorp/errwrap v1.1.0 // indirect
88+
github.com/hashicorp/go-multierror v1.1.1 // indirect
89+
github.com/jinzhu/inflection v1.0.0 // indirect
90+
github.com/jinzhu/now v1.1.5 // indirect
8291
github.com/jmespath/go-jmespath v0.4.1-0.20220621161143-b0104c826a24 // indirect
8392
github.com/josharian/intern v1.0.0 // indirect
8493
github.com/json-iterator/go v1.1.12 // indirect
@@ -107,13 +116,14 @@ require (
107116
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.59.0 // indirect
108117
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect
109118
go.opentelemetry.io/otel v1.35.0 // indirect
110-
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.28.0 // indirect
119+
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect
111120
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.27.0 // indirect
112121
go.opentelemetry.io/otel/metric v1.35.0 // indirect
113122
go.opentelemetry.io/otel/sdk v1.35.0 // indirect
114123
go.opentelemetry.io/otel/sdk/metric v1.35.0 // indirect
115124
go.opentelemetry.io/otel/trace v1.35.0 // indirect
116125
go.opentelemetry.io/proto/otlp v1.3.1 // indirect
126+
go.uber.org/atomic v1.9.0 // indirect
117127
go.uber.org/automaxprocs v1.6.0 // indirect
118128
golang.org/x/crypto v0.37.0 // indirect
119129
golang.org/x/oauth2 v0.29.0 // indirect

go.sum

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
8686
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
8787
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
8888
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
89+
github.com/dhui/dktest v0.4.5 h1:uUfYBIVREmj/Rw6MvgmqNAYzTiKOHJak+enB5Di73MM=
90+
github.com/dhui/dktest v0.4.5/go.mod h1:tmcyeHDKagvlDrz7gDKq4UAJOLIfVZYkfD5OnHDwcCo=
8991
github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
9092
github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
9193
github.com/docker/docker v28.0.1+incompatible h1:FCHjSRdXhNRFjlHMTv4jUNlIBbTeRjrWfeFuJp7jpo0=
@@ -153,13 +155,18 @@ github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/o
153155
github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY=
154156
github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY=
155157
github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
158+
github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
159+
github.com/go-sql-driver/mysql v1.7.1 h1:lUIinVbN1DY0xBg0eMOzmmtGoHwWBbvnWubQUrtU8EI=
160+
github.com/go-sql-driver/mysql v1.7.1/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
156161
github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0/go.mod h1:fyg7847qk6SyHyPtNmDHnmrv/HOrqktSC+C9fM+CJOE=
157162
github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI=
158163
github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
159164
github.com/go-viper/mapstructure/v2 v2.2.1 h1:ZAaOCxANMuZx5RCeg0mBdEZk7DZasvvZIxtHqx8aGss=
160165
github.com/go-viper/mapstructure/v2 v2.2.1/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
161166
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
162167
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
168+
github.com/golang-migrate/migrate/v4 v4.18.3 h1:EYGkoOsvgHHfm5U/naS1RP/6PL/Xv3S4B/swMiAmDLs=
169+
github.com/golang-migrate/migrate/v4 v4.18.3/go.mod h1:99BKpIi6ruaaXRM1A77eqZ+FWPQ3cfRa+ZVy5bmWMaY=
163170
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
164171
github.com/golang/glog v1.2.5 h1:DrW6hGnjIhtvhOIiAKT6Psh/Kd/ldepEa81DKeiRJ5I=
165172
github.com/golang/glog v1.2.5/go.mod h1:6AhwSGph0fcJtXVM/PEHPqZlFeoLxhs7/t5UDAwmO+w=
@@ -228,11 +235,20 @@ github.com/googleapis/google-cloud-go-testing v0.0.0-20210719221736-1c9a4c676720
228235
github.com/googleapis/google-cloud-go-testing v0.0.0-20210719221736-1c9a4c676720/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g=
229236
github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 h1:e9Rjr40Z98/clHv5Yg79Is0NtosR5LXRvdr7o/6NwbA=
230237
github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1/go.mod h1:tIxuGz/9mpox++sgp9fJjHO0+q1X9/UOWd798aAm22M=
238+
github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
239+
github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I=
240+
github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
241+
github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo=
242+
github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM=
231243
github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
232244
github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
233245
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
234246
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
235247
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
248+
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
249+
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
250+
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
251+
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
236252
github.com/jmespath/go-jmespath v0.4.1-0.20220621161143-b0104c826a24 h1:liMMTbpW34dhU4az1GN0pTPADwNmvoRSeoZ6PItiqnY=
237253
github.com/jmespath/go-jmespath v0.4.1-0.20220621161143-b0104c826a24/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
238254
github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
@@ -256,6 +272,8 @@ github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0
256272
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
257273
github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
258274
github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
275+
github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
276+
github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
259277
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4=
260278
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
261279
github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE=
@@ -367,6 +385,8 @@ github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8
367385
github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
368386
github.com/testcontainers/testcontainers-go v0.37.0 h1:L2Qc0vkTw2EHWQ08djon0D2uw7Z/PtHS/QzZZ5Ra/hg=
369387
github.com/testcontainers/testcontainers-go v0.37.0/go.mod h1:QPzbxZhQ6Bclip9igjLFj6z0hs01bU8lrl2dHQmgFGM=
388+
github.com/testcontainers/testcontainers-go/modules/mysql v0.37.0 h1:LqUos1oR5iuuzorFnSvxsHNdYdCHB/DfI82CuT58wbI=
389+
github.com/testcontainers/testcontainers-go/modules/mysql v0.37.0/go.mod h1:vHEEHx5Kf+uq5hveaVAMrTzPY8eeRZcKcl23MRw5Tkc=
370390
github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=
371391
github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=
372392
github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=
@@ -393,8 +413,8 @@ go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRND
393413
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ=
394414
go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ=
395415
go.opentelemetry.io/otel v1.35.0/go.mod h1:UEqy8Zp11hpkUrL73gSlELM0DupHoiq72dR+Zqel/+Y=
396-
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.28.0 h1:3Q/xZUyC1BBkualc9ROb4G8qkH90LXEIICcs5zv1OYY=
397-
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.28.0/go.mod h1:s75jGIWA9OfCMzF0xr+ZgfrB5FEbbV7UuYo32ahUiFI=
416+
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 h1:dIIDULZJpgdiHz5tXrTgKIMLkus6jEFa7x5SOKcyR7E=
417+
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0/go.mod h1:jlRVBe7+Z1wyxFSUs48L6OBQZ5JwH2Hg/Vbl+t9rAgI=
398418
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.27.0 h1:qFffATk0X+HD+f1Z8lswGiOQYKHRlzfmdJm0wEaVrFA=
399419
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.27.0/go.mod h1:MOiCmryaYtc+V0Ei+Tx9o5S1ZjA7kzLucuVuyzBZloQ=
400420
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.19.0 h1:IeMeyr1aBvBiPVYihXIaeIZba6b8E1bYp7lbdxK8CQg=
@@ -412,6 +432,8 @@ go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J
412432
go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0=
413433
go.opentelemetry.io/proto/otlp v1.3.1/go.mod h1:0X1WI4de4ZsLrrJNLAQbFeLCm3T7yBkR0XqQ7niQU+8=
414434
go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
435+
go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
436+
go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
415437
go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
416438
go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
417439
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
@@ -599,6 +621,10 @@ gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
599621
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
600622
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
601623
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
624+
gorm.io/driver/mysql v1.5.7 h1:MndhOPYOfEp2rHKgkZIhJ16eVUIRf2HmzgoPmh7FCWo=
625+
gorm.io/driver/mysql v1.5.7/go.mod h1:sEtPWMiqiN1N1cMXoXmBbd8C6/l+TESwriotuRRpkDM=
626+
gorm.io/gorm v1.25.7 h1:VsD6acwRjz2zFxGO50gPO6AkNs7KKnvfzUjHQhZDz/A=
627+
gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=
602628
gotest.tools/v3 v3.5.1 h1:EENdUnS3pdur5nybKYIh2Vfgc8IUNBjxDPSjtiJcOzU=
603629
gotest.tools/v3 v3.5.1/go.mod h1:isy3WKz7GK6uNw/sbHzfKBLvlvXwUyV06n6brMxxopU=
604630
honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

internal/datastore/connector.go

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import (
44
"errors"
55
"fmt"
66

7+
"github.com/kubeflow/model-registry/internal/datastore/embedmd"
8+
"github.com/kubeflow/model-registry/internal/datastore/mlmd"
79
"github.com/kubeflow/model-registry/pkg/api"
810
)
911

@@ -15,8 +17,9 @@ var (
1517
type TeardownFunc func() error
1618

1719
type Datastore struct {
18-
MLMD MLMDConfig
19-
Type string
20+
MLMD mlmd.MLMDConfig
21+
EmbedMD embedmd.EmbedMDConfig
22+
Type string
2023
}
2124

2225
type Connector interface {
@@ -31,10 +34,21 @@ func NewConnector(ds Datastore) (Connector, error) {
3134
return nil, fmt.Errorf("invalid MLMD config: %w", err)
3235
}
3336

34-
mlmd := NewMLMDService(&ds.MLMD)
37+
mlmd := mlmd.NewMLMDService(&ds.MLMD)
3538

3639
return mlmd, nil
40+
case "embedmd":
41+
if err := ds.EmbedMD.Validate(); err != nil {
42+
return nil, fmt.Errorf("invalid EmbedMD config: %w", err)
43+
}
44+
45+
embedmd, err := embedmd.NewEmbedMDService(&ds.EmbedMD)
46+
if err != nil {
47+
return nil, fmt.Errorf("error creating EmbedMD service: %w", err)
48+
}
49+
50+
return embedmd, nil
3751
default:
38-
return nil, fmt.Errorf("%w: %s. Supported types: mlmd", ErrUnsupportedDatastore, ds.Type)
52+
return nil, fmt.Errorf("%w: %s. Supported types: mlmd, embedmd", ErrUnsupportedDatastore, ds.Type)
3953
}
4054
}
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
package mysql
2+
3+
import (
4+
"gorm.io/driver/mysql"
5+
"gorm.io/gorm"
6+
)
7+
8+
type MySQLDBConnector struct {
9+
DSN string
10+
db *gorm.DB
11+
}
12+
13+
func NewMySQLDBConnector(dsn string) *MySQLDBConnector {
14+
return &MySQLDBConnector{
15+
DSN: dsn,
16+
}
17+
}
18+
19+
func (c *MySQLDBConnector) Connect() (*gorm.DB, error) {
20+
db, err := gorm.Open(mysql.Open(c.DSN), &gorm.Config{})
21+
if err != nil {
22+
return nil, err
23+
}
24+
25+
c.db = db
26+
27+
return db, nil
28+
}
29+
30+
func (c *MySQLDBConnector) DB() *gorm.DB {
31+
return c.db
32+
}
Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package mysql
2+
3+
import (
4+
"embed"
5+
"fmt"
6+
7+
"github.com/golang-migrate/migrate/v4"
8+
"github.com/golang-migrate/migrate/v4/database/mysql"
9+
"github.com/golang-migrate/migrate/v4/source/iofs"
10+
"gorm.io/gorm"
11+
)
12+
13+
//go:embed migrations/*.sql
14+
var migrations embed.FS
15+
16+
const (
17+
MigrationDir = "migrations"
18+
)
19+
20+
type MySQLMigrator struct {
21+
migrator *migrate.Migrate
22+
}
23+
24+
func NewMySQLMigrator(db *gorm.DB) (*MySQLMigrator, error) {
25+
sqlDB, err := db.DB()
26+
if err != nil {
27+
return nil, err
28+
}
29+
30+
driver, err := mysql.WithInstance(sqlDB, &mysql.Config{})
31+
if err != nil {
32+
return nil, err
33+
}
34+
35+
// Create a new source instance from the embedded files
36+
source, err := iofs.New(migrations, MigrationDir)
37+
if err != nil {
38+
return nil, err
39+
}
40+
41+
m, err := migrate.NewWithInstance(
42+
"iofs",
43+
source,
44+
"mysql",
45+
driver,
46+
)
47+
if err != nil {
48+
return nil, err
49+
}
50+
51+
return &MySQLMigrator{
52+
migrator: m,
53+
}, nil
54+
}
55+
56+
func (m *MySQLMigrator) Migrate() error {
57+
if err := m.Up(nil); err != nil && err != migrate.ErrNoChange {
58+
return err
59+
}
60+
61+
return nil
62+
}
63+
64+
func (m *MySQLMigrator) Up(steps *int) error {
65+
if steps == nil {
66+
return m.migrator.Up()
67+
}
68+
69+
if *steps < 0 {
70+
return fmt.Errorf("steps cannot be negative")
71+
}
72+
73+
return m.migrator.Steps(*steps)
74+
}
75+
76+
func (m *MySQLMigrator) Down(steps *int) error {
77+
if steps == nil {
78+
return m.migrator.Down()
79+
}
80+
81+
if *steps > 0 {
82+
return fmt.Errorf("steps cannot be positive")
83+
}
84+
85+
return m.migrator.Steps(*steps)
86+
}

0 commit comments

Comments
 (0)