From 661d55501bb5f9bd7f758ce3f2eeb8758a326c76 Mon Sep 17 00:00:00 2001 From: Liquan Pei Date: Mon, 4 Dec 2023 11:42:14 -0800 Subject: [PATCH] [BUG] Postgres support for metadata store (#1459) ## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - This PR fixes the dbcore code that originally connects MySQL family - New functionality - ... ## Test plan *How are these changes tested?* - [ ] Manual testing by connecting to Postgres ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* --- go/coordinator/cmd/grpccoordinator/cmd.go | 3 ++- go/coordinator/internal/grpccoordinator/server.go | 2 ++ go/coordinator/internal/metastore/db/dbcore/core.go | 5 +++-- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/go/coordinator/cmd/grpccoordinator/cmd.go b/go/coordinator/cmd/grpccoordinator/cmd.go index d4dc50d0dd3..d3f76b0c6a4 100644 --- a/go/coordinator/cmd/grpccoordinator/cmd.go +++ b/go/coordinator/cmd/grpccoordinator/cmd.go @@ -26,7 +26,8 @@ func init() { Cmd.Flags().StringVar(&conf.SystemCatalogProvider, "system-catalog-provider", "memory", "System catalog provider") Cmd.Flags().StringVar(&conf.Username, "username", "root", "MetaTable username") Cmd.Flags().StringVar(&conf.Password, "password", "", "MetaTable password") - Cmd.Flags().StringVar(&conf.Address, "db-address", "127.0.0.1:3306", "MetaTable db address") + Cmd.Flags().StringVar(&conf.Address, "db-address", "127.0.0.1", "MetaTable db address") + Cmd.Flags().IntVar(&conf.Port, "db-port", 5432, "MetaTable db port") Cmd.Flags().StringVar(&conf.DBName, "db-name", "", "MetaTable db name") Cmd.Flags().IntVar(&conf.MaxIdleConns, "max-idle-conns", 10, "MetaTable max idle connections") Cmd.Flags().IntVar(&conf.MaxOpenConns, "max-open-conns", 10, "MetaTable max open connections") diff --git a/go/coordinator/internal/grpccoordinator/server.go b/go/coordinator/internal/grpccoordinator/server.go index ac99a18bdf0..fc4eb850b29 100644 --- a/go/coordinator/internal/grpccoordinator/server.go +++ b/go/coordinator/internal/grpccoordinator/server.go @@ -28,6 +28,7 @@ type Config struct { Username string Password string Address string + Port int DBName string MaxIdleConns int MaxOpenConns int @@ -69,6 +70,7 @@ func New(config Config) (*Server, error) { Username: config.Username, Password: config.Password, Address: config.Address, + Port: config.Port, DBName: config.DBName, MaxIdleConns: config.MaxIdleConns, MaxOpenConns: config.MaxOpenConns, diff --git a/go/coordinator/internal/metastore/db/dbcore/core.go b/go/coordinator/internal/metastore/db/dbcore/core.go index 22ebc8c0597..253a63224c0 100644 --- a/go/coordinator/internal/metastore/db/dbcore/core.go +++ b/go/coordinator/internal/metastore/db/dbcore/core.go @@ -24,14 +24,15 @@ type DBConfig struct { Username string Password string Address string + Port int DBName string MaxIdleConns int MaxOpenConns int } func Connect(cfg DBConfig) (*gorm.DB, error) { - dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8mb4&parseTime=True&loc=Local&tls=true&interpolateParams=true", - cfg.Username, cfg.Password, cfg.Address, cfg.DBName) + dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%d", + cfg.Address, cfg.Username, cfg.Password, cfg.DBName, cfg.Port) ormLogger := logger.Default ormLogger.LogMode(logger.Info)