From 9bef06f47fff6aa13b862e046bd7e70057115d15 Mon Sep 17 00:00:00 2001 From: Kartikay Date: Sun, 9 Mar 2025 05:18:02 +0530 Subject: [PATCH] granular conn flags for datastores Signed-off-by: Kartikay --- pkg/cmd/datastore/datastore.go | 78 ++++++++++++++++++++++ pkg/cmd/datastore/zz_generated.options.go | 81 +++++++++++++++++++++++ 2 files changed, 159 insertions(+) diff --git a/pkg/cmd/datastore/datastore.go b/pkg/cmd/datastore/datastore.go index 256c8855e..088c4f610 100644 --- a/pkg/cmd/datastore/datastore.go +++ b/pkg/cmd/datastore/datastore.go @@ -102,6 +102,12 @@ func deprecateUnifiedConnFlags(flagSet *pflag.FlagSet) { type Config struct { Engine string `debugmap:"visible"` URI string `debugmap:"sensitive"` + DatastoreHost string `debugmap:"sensitive"` + DatastorePort string `debugmap:"sensitive"` + DatastoreUsername string `debugmap:"sensitive"` + DatastorePassword string `debugmap:"sensitive"` + DatastoreName string `debugmap:"sensitive"` + DatastoreSSLMode string `debugmap:"sensitive"` GCWindow time.Duration `debugmap:"visible"` LegacyFuzzing time.Duration `debugmap:"visible"` RevisionQuantization time.Duration `debugmap:"visible"` @@ -153,6 +159,9 @@ type Config struct { SpannerMinSessions uint64 `debugmap:"visible"` SpannerMaxSessions uint64 `debugmap:"visible"` SpannerDatastoreMetricsOption string `debugmap:"visible"` + SpannerInstanceID string `debugmap:"sensitive"` + SpannerDatabaseID string `debugmap:"sensitive"` + SpannerProjectID string `debugmap:"sensitive"` // MySQL TablePrefix string `debugmap:"visible"` @@ -283,12 +292,28 @@ func RegisterDatastoreFlagsWithPrefix(flagSet *pflag.FlagSet, prefix string, opt flagSet.BoolVar(&opts.ExperimentalColumnOptimization, flagName("datastore-experimental-column-optimization"), false, "enable experimental column optimization") + flagSet.StringVar(&opts.DatastoreHost, flagName("datastore-host"), defaults.DatastoreHost, "database host for datastore connection string") + flagSet.StringVar(&opts.DatastorePort, flagName("datastore-port"), defaults.DatastorePort, "database port for datastore connection string") + flagSet.StringVar(&opts.DatastoreUsername, flagName("datastore-user"), defaults.DatastoreUsername, "database username for datastore connection string") + flagSet.StringVar(&opts.DatastorePassword, flagName("datastore-password"), defaults.DatastorePassword, "database password for datastore connection string") + flagSet.StringVar(&opts.DatastoreName, flagName("datastore-name"), defaults.DatastoreName, "database name for datastore connection string") + flagSet.StringVar(&opts.DatastoreSSLMode, flagName("datastore-ssl-mode"), defaults.DatastoreSSLMode, "database SSL mode for datastore connection string") + flagSet.StringVar(&opts.SpannerProjectID, flagName("spanner-datastore-project-id"), defaults.SpannerProjectID, "spanner project-id for datastore connection string") + flagSet.StringVar(&opts.SpannerInstanceID, flagName("spanner-datastore-instance-id"), defaults.SpannerInstanceID, "spanner instance-id for datastore connection string") + flagSet.StringVar(&opts.SpannerDatabaseID, flagName("spanner-datastore-database-id"), defaults.SpannerDatabaseID, "spanner database-id for datastore connection string") + return nil } func DefaultDatastoreConfig() *Config { return &Config{ Engine: MemoryEngine, + DatastoreHost: "", + DatastorePort: "", + DatastoreUsername: "", + DatastorePassword: "", + DatastoreName: "", + DatastoreSSLMode: "", GCWindow: 24 * time.Hour, LegacyFuzzing: -1, RevisionQuantization: 5 * time.Second, @@ -319,6 +344,9 @@ func DefaultDatastoreConfig() *Config { RequestHedgingQuantile: 0.95, SpannerCredentialsFile: "", SpannerEmulatorHost: "", + SpannerInstanceID: "", + SpannerDatabaseID: "", + SpannerProjectID: "", TablePrefix: "", MigrationPhase: "", FollowerReadDelay: DefaultFollowerReadDelay, @@ -343,6 +371,56 @@ func NewDatastore(ctx context.Context, options ...ConfigOption) (datastore.Datas o(opts) } + if opts.URI == "" && ((opts.DatastoreHost != "" && opts.DatastoreUsername != "" && opts.DatastoreName != "") || (opts.SpannerDatabaseID != "" && opts.SpannerInstanceID != "" && opts.SpannerProjectID != "")) { + switch opts.Engine { + case PostgresEngine, CockroachEngine: + sslMode := opts.DatastoreSSLMode + if sslMode == "" { + sslMode = "prefer" + } + + port := opts.DatastorePort + if port == "" { + if opts.Engine == PostgresEngine { + port = "5432" + } else { + port = "26257" + } + } + + userPart := "" + if opts.DatastoreUsername != "" { + userPart = opts.DatastoreUsername + if opts.DatastorePassword != "" { + userPart = userPart + ":" + opts.DatastorePassword + } + userPart = userPart + "@" + } + + opts.URI = fmt.Sprintf("%s://%s%s:%s/%s?sslmode=%s", + opts.Engine, userPart, opts.DatastoreHost, port, opts.DatastoreName, sslMode) + + case MySQLEngine: + port := opts.DatastorePort + if port == "" { + port = "3306" + } + opts.URI = fmt.Sprintf("%s:%s@(%s:%s)/%s?parseTime=True", + opts.DatastoreUsername, opts.DatastorePassword, opts.DatastoreHost, port, opts.DatastoreName) + + case SpannerEngine: + if opts.SpannerDatabaseID != "" && opts.SpannerInstanceID != "" && opts.SpannerProjectID != "" { + opts.URI = fmt.Sprintf("projects/%s/instances/%s/databases/%s", + opts.SpannerProjectID, opts.SpannerInstanceID, opts.SpannerDatabaseID) + } else { + return nil, fmt.Errorf("%s engine requires host (project) and database parameters", SpannerEngine) + } + + default: + return nil, fmt.Errorf("engine %s does not support connection parameter construction", opts.Engine) + } + } + if (opts.Engine == PostgresEngine || opts.Engine == MySQLEngine) && opts.FollowerReadDelay == DefaultFollowerReadDelay { // Set the default follower read delay for postgres and mysql to 0 - // this should only be set if read replicas are used. diff --git a/pkg/cmd/datastore/zz_generated.options.go b/pkg/cmd/datastore/zz_generated.options.go index 7fe1beab8..18789f3e3 100644 --- a/pkg/cmd/datastore/zz_generated.options.go +++ b/pkg/cmd/datastore/zz_generated.options.go @@ -33,6 +33,12 @@ func (c *Config) ToOption() ConfigOption { return func(to *Config) { to.Engine = c.Engine to.URI = c.URI + to.DatastoreHost = c.DatastoreHost + to.DatastorePort = c.DatastorePort + to.DatastoreUsername = c.DatastoreUsername + to.DatastorePassword = c.DatastorePassword + to.DatastoreName = c.DatastoreName + to.DatastoreSSLMode = c.DatastoreSSLMode to.GCWindow = c.GCWindow to.LegacyFuzzing = c.LegacyFuzzing to.RevisionQuantization = c.RevisionQuantization @@ -70,6 +76,9 @@ func (c *Config) ToOption() ConfigOption { to.SpannerMinSessions = c.SpannerMinSessions to.SpannerMaxSessions = c.SpannerMaxSessions to.SpannerDatastoreMetricsOption = c.SpannerDatastoreMetricsOption + to.SpannerInstanceID = c.SpannerInstanceID + to.SpannerDatabaseID = c.SpannerDatabaseID + to.SpannerProjectID = c.SpannerProjectID to.TablePrefix = c.TablePrefix to.RelationshipIntegrityEnabled = c.RelationshipIntegrityEnabled to.RelationshipIntegrityCurrentKey = c.RelationshipIntegrityCurrentKey @@ -90,6 +99,12 @@ func (c Config) DebugMap() map[string]any { debugMap := map[string]any{} debugMap["Engine"] = helpers.DebugValue(c.Engine, false) debugMap["URI"] = helpers.SensitiveDebugValue(c.URI) + debugMap["DatastoreHost"] = helpers.SensitiveDebugValue(c.DatastoreHost) + debugMap["DatastorePort"] = helpers.SensitiveDebugValue(c.DatastorePort) + debugMap["DatastoreUsername"] = helpers.SensitiveDebugValue(c.DatastoreUsername) + debugMap["DatastorePassword"] = helpers.SensitiveDebugValue(c.DatastorePassword) + debugMap["DatastoreName"] = helpers.SensitiveDebugValue(c.DatastoreName) + debugMap["DatastoreSSLMode"] = helpers.SensitiveDebugValue(c.DatastoreSSLMode) debugMap["GCWindow"] = helpers.DebugValue(c.GCWindow, false) debugMap["LegacyFuzzing"] = helpers.DebugValue(c.LegacyFuzzing, false) debugMap["RevisionQuantization"] = helpers.DebugValue(c.RevisionQuantization, false) @@ -126,6 +141,9 @@ func (c Config) DebugMap() map[string]any { debugMap["SpannerMinSessions"] = helpers.DebugValue(c.SpannerMinSessions, false) debugMap["SpannerMaxSessions"] = helpers.DebugValue(c.SpannerMaxSessions, false) debugMap["SpannerDatastoreMetricsOption"] = helpers.DebugValue(c.SpannerDatastoreMetricsOption, false) + debugMap["SpannerInstanceID"] = helpers.SensitiveDebugValue(c.SpannerInstanceID) + debugMap["SpannerDatabaseID"] = helpers.SensitiveDebugValue(c.SpannerDatabaseID) + debugMap["SpannerProjectID"] = helpers.SensitiveDebugValue(c.SpannerProjectID) debugMap["TablePrefix"] = helpers.DebugValue(c.TablePrefix, false) debugMap["RelationshipIntegrityEnabled"] = helpers.DebugValue(c.RelationshipIntegrityEnabled, false) debugMap["RelationshipIntegrityCurrentKey"] = helpers.DebugValue(c.RelationshipIntegrityCurrentKey, false) @@ -171,6 +189,48 @@ func WithURI(uRI string) ConfigOption { } } +// WithDatastoreHost returns an option that can set DatastoreHost on a Config +func WithDatastoreHost(datastoreHost string) ConfigOption { + return func(c *Config) { + c.DatastoreHost = datastoreHost + } +} + +// WithDatastorePort returns an option that can set DatastorePort on a Config +func WithDatastorePort(datastorePort string) ConfigOption { + return func(c *Config) { + c.DatastorePort = datastorePort + } +} + +// WithDatastoreUsername returns an option that can set DatastoreUsername on a Config +func WithDatastoreUsername(datastoreUsername string) ConfigOption { + return func(c *Config) { + c.DatastoreUsername = datastoreUsername + } +} + +// WithDatastorePassword returns an option that can set DatastorePassword on a Config +func WithDatastorePassword(datastorePassword string) ConfigOption { + return func(c *Config) { + c.DatastorePassword = datastorePassword + } +} + +// WithDatastoreName returns an option that can set DatastoreName on a Config +func WithDatastoreName(datastoreName string) ConfigOption { + return func(c *Config) { + c.DatastoreName = datastoreName + } +} + +// WithDatastoreSSLMode returns an option that can set DatastoreSSLMode on a Config +func WithDatastoreSSLMode(datastoreSSLMode string) ConfigOption { + return func(c *Config) { + c.DatastoreSSLMode = datastoreSSLMode + } +} + // WithGCWindow returns an option that can set GCWindow on a Config func WithGCWindow(gCWindow time.Duration) ConfigOption { return func(c *Config) { @@ -458,6 +518,27 @@ func WithSpannerDatastoreMetricsOption(spannerDatastoreMetricsOption string) Con } } +// WithSpannerInstanceID returns an option that can set SpannerInstanceID on a Config +func WithSpannerInstanceID(spannerInstanceID string) ConfigOption { + return func(c *Config) { + c.SpannerInstanceID = spannerInstanceID + } +} + +// WithSpannerDatabaseID returns an option that can set SpannerDatabaseID on a Config +func WithSpannerDatabaseID(spannerDatabaseID string) ConfigOption { + return func(c *Config) { + c.SpannerDatabaseID = spannerDatabaseID + } +} + +// WithSpannerProjectID returns an option that can set SpannerProjectID on a Config +func WithSpannerProjectID(spannerProjectID string) ConfigOption { + return func(c *Config) { + c.SpannerProjectID = spannerProjectID + } +} + // WithTablePrefix returns an option that can set TablePrefix on a Config func WithTablePrefix(tablePrefix string) ConfigOption { return func(c *Config) {