testing config

configuration.go

@@ -35,6 +35,8 @@ var (
 	defaultDatabaseConfig     = &configuration{flag: "default-database", envFlag: "default_database", defaultValue: ""}
 	defaultTableConfig        = &configuration{flag: "default-table", envFlag: "default_table", defaultValue: ""}
 	enableSigV4AuthConfig     = &configuration{flag: "enable-sigv4-auth", envFlag: "enable_sigv4_auth", defaultValue: "true"}
+	readBaseEndpointConfig    = &configuration{flag: "read-base-endpoint", envFlag: "read_base_endpoint", defaultValue: ""}
+	writeBaseEndpointConfig   = &configuration{flag: "write-base-endpoint", envFlag: "write_base_endpoint", defaultValue: ""}
 	listenAddrConfig          = &configuration{flag: "web.listen-address", envFlag: "", defaultValue: ":9201"}
 	telemetryPathConfig       = &configuration{flag: "web.telemetry-path", envFlag: "", defaultValue: "/metrics"}
 	failOnLabelConfig         = &configuration{flag: "fail-on-long-label", envFlag: "fail_on_long_label", defaultValue: "false"}

main.go

@@ -103,6 +103,8 @@ type connectionConfig struct {
 	telemetryPath             string
 	maxReadRetries            int
 	maxWriteRetries           int
+	readBaseEndpoint          string
+	writeBaseEndpoint         string
 	certificate               string
 	key                       string
 }
@@ -122,13 +124,13 @@ func main() {
 		logger := cfg.createLogger()
 
 		ctx := context.Background()
-		awsQueryConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxReadRetries)
+		awsQueryConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxReadRetries, cfg.readBaseEndpoint)
 		if err != nil {
 			timestream.LogError(logger, "Failed to build AWS configuration for query", err)
 			os.Exit(1)
 		}
 
-		awsWriteConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxWriteRetries)
+		awsWriteConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxWriteRetries, cfg.writeBaseEndpoint)
 		if err != nil {
 			timestream.LogError(logger, "Failed to build AWS configuration for write", err)
 			os.Exit(1)
@@ -185,12 +187,12 @@ func lambdaHandler(req events.APIGatewayProxyRequest) (events.APIGatewayProxyRes
 			return createErrorResponse(errors.NewParseBasicAuthHeaderError().(*errors.ParseBasicAuthHeaderError).Message())
 		}
 	}
-	awsQueryConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxReadRetries)
+	awsQueryConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxReadRetries, cfg.readBaseEndpoint)
 	if err != nil {
 		timestream.LogError(logger, "Failed to build AWS configuration for query", err)
 		os.Exit(1)
 	}
-	awsWriteConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxWriteRetries)
+	awsWriteConfigs, err := cfg.buildAWSConfig(ctx, cfg.maxWriteRetries, cfg.writeBaseEndpoint)
 	if err != nil {
 		timestream.LogError(logger, "Failed to build AWS configuration for write", err)
 		os.Exit(1)
@@ -394,6 +396,9 @@ func parseEnvironmentVariables() (*connectionConfig, error) {
 		return nil, errors.NewParseRetriesError(writeRetries, "write")
 	}
 
+	cfg.readBaseEndpoint = getOrDefault(readBaseEndpointConfig)
+	cfg.writeBaseEndpoint = getOrDefault(writeBaseEndpointConfig)
+
 	cfg.promlogConfig = promlog.Config{Level: &promlog.AllowedLevel{}, Format: &promlog.AllowedFormat{}}
 	cfg.promlogConfig.Level.Set(getOrDefault(promlogLevelConfig))
 	cfg.promlogConfig.Format.Set(getOrDefault(promlogFormatConfig))
@@ -415,6 +420,8 @@ func parseFlags() *connectionConfig {
 	var enableSigV4Auth string
 	var failOnLongMetricLabelName string
 	var failOnInvalidSample string
+	var readBaseEndpoint string
+	var writeBaseEndpoint string
 
 	a.Flag(enableLogConfig.flag, "Enables or disables logging in the connector. Default to 'true'.").Default(enableLogConfig.defaultValue).StringVar(&enableLogging)
 	a.Flag(regionConfig.flag, "The signing region for the Timestream service. Default to 'us-east-1'.").Default(regionConfig.defaultValue).StringVar(&cfg.clientConfig.region)
@@ -431,6 +438,10 @@ func parseFlags() *connectionConfig {
 	a.Flag(certificateConfig.flag, "TLS server certificate file.").Default(certificateConfig.defaultValue).StringVar(&cfg.certificate)
 	a.Flag(keyConfig.flag, "TLS server private key file.").Default(keyConfig.defaultValue).StringVar(&cfg.key)
 	a.Flag(enableSigV4AuthConfig.flag, "Whether to enable SigV4 authentication with the API Gateway. Default to 'false'.").Default(enableSigV4AuthConfig.defaultValue).StringVar(&enableSigV4Auth)
+	a.Flag(readBaseEndpointConfig.flag, "Override the default Timestream query endpoint (e.g., a VPC Endpoint).").
+		Default(readBaseEndpointConfig.defaultValue).StringVar(&readBaseEndpoint)
+	a.Flag(writeBaseEndpointConfig.flag, "Override the default Timestream write endpoint (e.g., a VPC Endpoint).").
+		Default(writeBaseEndpointConfig.defaultValue).StringVar(&writeBaseEndpoint)
 
 	flag.AddFlags(a, &cfg.promlogConfig)
 
@@ -457,7 +468,7 @@ func parseFlags() *connectionConfig {
 }
 
 // buildAWSConfig builds a aws.Config and return the pointer of the config.
-func (cfg *connectionConfig) buildAWSConfig(ctx context.Context, maxRetries int) (aws.Config, error) {
+func (cfg *connectionConfig) buildAWSConfig(ctx context.Context, maxRetries int, baseEndpoint string) (aws.Config, error) {
 	awsConfig, err := config.LoadDefaultConfig(ctx,
 		config.WithRegion(cfg.clientConfig.region),
 		config.WithRetryer(func() aws.Retryer {
@@ -469,6 +480,9 @@ func (cfg *connectionConfig) buildAWSConfig(ctx context.Context, maxRetries int)
 	if err != nil {
 		return aws.Config{}, fmt.Errorf("failed to build AWS config: %w", err)
 	}
+	if baseEndpoint != "" {
+		awsConfig.BaseEndpoint = aws.String(baseEndpoint)
+	}
 	return awsConfig, nil
 }
 

private-link.yaml

@@ -0,0 +1,22 @@
+services:
+  timestream-prometheus-connector:
+    container_name: connector
+    build: .
+    ports:
+      - "9201:9201"
+    volumes:
+      - .:/home
+    command:
+      - --default-database=${DEFAULT_DATABASE:-DevPrometheusDatabase}
+      - --default-table=${DEFAULT_TABLE:-DevPrometheusMetricsTable}
+      - --region=${AWS_REGION:-us-west-2}
+      - --log.level=debug
+    environment:
+      read_base_endpoint: https://query-cell1.timestream.us-west-2.amazonaws.com
+      write_base_endpoint: https://ingest-cell1.timestream.us-west-2.amazonaws.com
+    networks:
+      - aws_network
+
+networks:
+  aws_network:
+    external: true