Skip to content

Commit

Permalink
(feat) add source-access-token functionality
Browse files Browse the repository at this point in the history
  • Loading branch information
@tphoney
tphoney committed Oct 31, 2024
1 parent 0fa8789 commit 4e2a033
Show file tree
Hide file tree
Showing 5 changed files with 133 additions and 49 deletions.
47 changes: 37 additions & 10 deletions cmd.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"runtime"

"github.com/google/uuid"
"github.com/overmindtech/sdp-go"
"github.com/spf13/cobra"
"github.com/spf13/viper"
)
Expand All @@ -19,11 +20,27 @@ func AddEngineFlags(command *cobra.Command) {
cobra.CheckErr(viper.BindEnv("app", "APP"))
command.PersistentFlags().String("api-key", "", "The API key to use to authenticate to the Overmind API")
cobra.CheckErr(viper.BindEnv("api-key", "OVM_API_KEY", "API_KEY"))
command.PersistentFlags().String("source-access-token", "", "The access token to use to authenticate the source for managed sources")
cobra.CheckErr(viper.BindEnv("source-access-token", "SOURCE_ACCESS_TOKEN"))
// SOURCE_TOKEN_TYPE
command.PersistentFlags().String("source-token-type", "", "The type of token to use to authenticate the source for managed sources")
cobra.CheckErr(viper.BindEnv("source-token-type", "SOURCE_TOKEN_TYPE"))
command.PersistentFlags().Int("max-parallel", 0, "The maximum number of parallel executions")
cobra.CheckErr(viper.BindEnv("max-parallel", "MAX_PARALLEL"))
}

func EngineConfigFromViper(engineType, version string) (*EngineConfig, error) {
var sourceName string
if viper.GetString("source-name") == "" {
hostname, err := os.Hostname()
if err != nil {
return nil, fmt.Errorf("error getting hostname: %w", err)
}
sourceName = fmt.Sprintf("%s-%s", engineType, hostname)
} else {
sourceName = viper.GetString("source-name")
}

sourceUUIDString := viper.GetString("source-uuid")
var sourceUUID uuid.UUID
if sourceUUIDString == "" {
Expand All @@ -35,25 +52,35 @@ func EngineConfigFromViper(engineType, version string) (*EngineConfig, error) {
return nil, fmt.Errorf("error parsing source-uuid: %w", err)
}
}

// checks for source-access-token and api-key, if both set error, if none set error
if viper.GetString("source-access-token") != "" && viper.GetString("api-key") != "" {
return nil, fmt.Errorf("source-access-token and api-key cannot be set at the same time")
}
if viper.GetString("source-access-token") == "" && viper.GetString("api-key") == "" {
return nil, fmt.Errorf("source-access-token or api-key must be set")
}

var managedSource sdp.SourceManaged
if viper.GetString("source-access-token") != "" {
managedSource = sdp.SourceManaged_MANAGED
} else {
managedSource = sdp.SourceManaged_LOCAL
}

maxParallelExecutions := viper.GetInt("max-parallel")
if maxParallelExecutions == 0 {
maxParallelExecutions = runtime.NumCPU()
}
var sourceName string
if viper.GetString("source-name") == "" {
hostname, err := os.Hostname()
if err != nil {
return nil, fmt.Errorf("error getting hostname: %w", err)
}
sourceName = fmt.Sprintf("%s-%s", engineType, hostname)
} else {
sourceName = viper.GetString("source-name")
}

return &EngineConfig{
EngineType: engineType,
Version: version,
SourceName: sourceName,
SourceUUID: sourceUUID,
SourceAccessToken: viper.GetString("source-access-token"),
SourceTokenType: viper.GetString("source-token-type"),
OvermindManagedSource: managedSource,
App: viper.GetString("app"),
ApiKey: viper.GetString("api-key"),
MaxParallelExecutions: maxParallelExecutions,
Expand Down
108 changes: 82 additions & 26 deletions cmd_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,38 +6,46 @@ import (
"testing"

"github.com/google/uuid"
"github.com/overmindtech/sdp-go"
"github.com/spf13/viper"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func TestEngineConfigFromViper(t *testing.T) {
tests := []struct {
name string
setupViper func()
engineType string
version string
expectedSourceName string
expectedSourceUUID uuid.UUID
expectedApp string
expectedApiKey string
expectedMaxParallel int
expectError bool
name string
setupViper func()
engineType string
version string
expectedSourceName string
expectedSourceUUID uuid.UUID
expectedSourceAccessToken string
expectedSourceTokenType string
expectedManagedSource sdp.SourceManaged
expectedApp string
expectedApiKey string
expectedMaxParallel int
expectError bool
}{
{
name: "default values",
setupViper: func() {
viper.Reset()
viper.Set("app", "https://app.overmind.tech")
viper.Set("api-key", "api-key")
},
engineType: "test-engine",
version: "1.0",
expectedSourceName: "test-engine-" + getHostname(t),
expectedSourceUUID: uuid.Nil,
expectedApp: "https://app.overmind.tech",
expectedApiKey: "",
expectedMaxParallel: runtime.NumCPU(),
expectError: false,
engineType: "test-engine",
version: "1.0",
expectedSourceName: "test-engine-" + getHostname(t),
expectedSourceUUID: uuid.Nil,
expectedSourceAccessToken: "",
expectedSourceTokenType: "",
expectedManagedSource: sdp.SourceManaged_LOCAL,
expectedApp: "https://app.overmind.tech",
expectedApiKey: "api-key",
expectedMaxParallel: runtime.NumCPU(),
expectError: false,
},
{
name: "custom values",
Expand All @@ -49,14 +57,17 @@ func TestEngineConfigFromViper(t *testing.T) {
viper.Set("api-key", "custom-api-key")
viper.Set("max-parallel", 10)
},
engineType: "test-engine",
version: "1.0",
expectedSourceName: "custom-source",
expectedSourceUUID: uuid.MustParse("123e4567-e89b-12d3-a456-426614174000"),
expectedApp: "https://custom.app",
expectedApiKey: "custom-api-key",
expectedMaxParallel: 10,
expectError: false,
engineType: "test-engine",
version: "1.0",
expectedSourceName: "custom-source",
expectedSourceUUID: uuid.MustParse("123e4567-e89b-12d3-a456-426614174000"),
expectedSourceAccessToken: "",
expectedSourceTokenType: "",
expectedManagedSource: sdp.SourceManaged_LOCAL,
expectedApp: "https://custom.app",
expectedApiKey: "custom-api-key",
expectedMaxParallel: 10,
expectError: false,
},
{
name: "invalid UUID",
Expand All @@ -68,6 +79,48 @@ func TestEngineConfigFromViper(t *testing.T) {
version: "1.0",
expectError: true,
},
{
name: "custom values with source access token",
setupViper: func() {
viper.Reset()
viper.Set("source-name", "custom-source")
viper.Set("source-uuid", "123e4567-e89b-12d3-a456-426614174000")
viper.Set("app", "https://custom.app")
viper.Set("source-access-token", "custom-access-token")
viper.Set("source-token-type", "custom-token-type")
viper.Set("max-parallel", 10)
},
engineType: "test-engine",
version: "1.0",
expectedSourceName: "custom-source",
expectedSourceUUID: uuid.MustParse("123e4567-e89b-12d3-a456-426614174000"),
expectedSourceAccessToken: "custom-access-token",
expectedSourceTokenType: "custom-token-type",
expectedManagedSource: sdp.SourceManaged_MANAGED,
expectedApp: "https://custom.app",
expectedMaxParallel: 10,
expectError: false,
},
{
name: "source access token and api key set",
setupViper: func() {
viper.Reset()
viper.Set("source-access-token", "custom-access-token")
viper.Set("api-key", "custom-api-key")
},
engineType: "test-engine",
version: "1.0",
expectError: true,
},
{
name: "source access token and api key not set",
setupViper: func() {
viper.Reset()
},
engineType: "test-engine",
version: "1.0",
expectError: true,
},
}

for _, tt := range tests {
Expand All @@ -86,6 +139,9 @@ func TestEngineConfigFromViper(t *testing.T) {
} else {
assert.Equal(t, tt.expectedSourceUUID, config.SourceUUID)
}
assert.Equal(t, tt.expectedSourceAccessToken, config.SourceAccessToken)
assert.Equal(t, tt.expectedSourceTokenType, config.SourceTokenType)
assert.Equal(t, tt.expectedManagedSource, config.OvermindManagedSource)
assert.Equal(t, tt.expectedApp, config.App)
assert.Equal(t, tt.expectedApiKey, config.ApiKey)
assert.Equal(t, tt.expectedMaxParallel, config.MaxParallelExecutions)
Expand Down
24 changes: 13 additions & 11 deletions engine.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,13 +43,19 @@ type HeartbeatOptions struct {
// EngineConfig is the configuration for the engine
// it is used to configure the engine before starting it
type EngineConfig struct {
EngineType string // The type of the engine, e.g. "aws" or "kubernetes"
Version string // The version of the adapter that should be reported in the heartbeat
SourceName string // normally follows the format of "type-hostname", e.g. "stdlib-source"
SourceUUID uuid.UUID // The UUID of the source, is this is blank it will be auto-generated. This is used in heartbeats and shouldn't be supplied usually"
App string // "https://app.overmind.tech", "The URL of the Overmind app to use"
ApiKey string // The API key to use to authenticate to the Overmind API"
MaxParallelExecutions int // 2_000, Max number of requests to run in parallel
EngineType string // The type of the engine, e.g. "aws" or "kubernetes"
Version string // The version of the adapter that should be reported in the heartbeat
SourceName string // normally follows the format of "type-hostname", e.g. "stdlib-source"
SourceUUID uuid.UUID // The UUID of the source, is this is blank it will be auto-generated. This is used in heartbeats and shouldn't be supplied usually"
App string // "https://app.overmind.tech", "The URL of the Overmind app to use"
ApiKey string // The API key to use to authenticate to the Overmind API"
SourceAccessToken string // The access token to use to authenticate to the source
SourceTokenType string // The type of token to use to authenticate the source for managed sources
// Whether this adapter is managed by Overmind. This is initially used for
// reporting so that you can tell the difference between managed adapters and
// ones you're running locally
OvermindManagedSource sdp.SourceManaged
MaxParallelExecutions int // 2_000, Max number of requests to run in parallel
}

// Engine is the main discovery engine. This is where all of the Adapters and
Expand All @@ -60,10 +66,6 @@ type EngineConfig struct {
// simply not communicate over NATS
type Engine struct {
EngineConfig *EngineConfig
// Whether this adapter is managed by Overmind. This is initially used for
// reporting so that you can tell the difference between managed adapters and
// ones you're running locally
Managed sdp.SourceManaged

NATSOptions *auth.NATSOptions // Options for connecting to NATS
NATSQueueName string // The name of the queue to use when subscribing
Expand Down
2 changes: 1 addition & 1 deletion heartbeat.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ func (e *Engine) SendHeartbeat(ctx context.Context) error {
Type: e.EngineConfig.EngineType,
AvailableScopes: availableScopes,
AdapterMetadata: adapterMetadata,
Managed: e.Managed,
Managed: e.EngineConfig.OvermindManagedSource,
Error: heartbeatError,
NextHeartbeatMax: durationpb.New(nextHeartbeat),
},
Expand Down
1 change: 0 additions & 1 deletion heartbeat_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ func TestHeartbeats(t *testing.T) {
}
e, _ := NewEngine(&ec)

e.Managed = sdp.SourceManaged_LOCAL
e.HeartbeatOptions = &HeartbeatOptions{
ManagementClient: testHeartbeatClient{
Requests: requests,
Expand Down

0 comments on commit 4e2a033

Please sign in to comment.