forked from flyteorg/flyte
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Gleb Kanterov <[email protected]>
- Loading branch information
Showing
17 changed files
with
1,714 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,19 @@ | ||
package google | ||
|
||
type TokenSourceFactoryType = string | ||
|
||
const ( | ||
TokenSourceTypeDefault = "default" | ||
) | ||
|
||
type TokenSourceFactoryConfig struct { | ||
// Type is type of TokenSourceFactory, possible values are 'default' or 'gke'. | ||
// - 'default' uses default credentials, see https://cloud.google.com/iam/docs/service-accounts#default | ||
Type TokenSourceFactoryType `json:"type" pflag:",Defines type of TokenSourceFactory, possible values are 'default'"` | ||
} | ||
|
||
func GetDefaultConfig() TokenSourceFactoryConfig { | ||
return TokenSourceFactoryConfig{ | ||
Type: "default", | ||
} | ||
} |
18 changes: 18 additions & 0 deletions
18
flyteplugins/go/tasks/pluginmachinery/google/default_token_source_factory.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
package google | ||
|
||
import ( | ||
"context" | ||
|
||
"golang.org/x/oauth2" | ||
"golang.org/x/oauth2/google" | ||
) | ||
|
||
type defaultTokenSource struct{} | ||
|
||
func (m *defaultTokenSource) GetTokenSource(ctx context.Context, identity Identity) (oauth2.TokenSource, error) { | ||
return google.DefaultTokenSource(ctx) | ||
} | ||
|
||
func NewDefaultTokenSourceFactory() (TokenSourceFactory, error) { | ||
return &defaultTokenSource{}, nil | ||
} |
25 changes: 25 additions & 0 deletions
25
flyteplugins/go/tasks/pluginmachinery/google/token_source_factory.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
package google | ||
|
||
import ( | ||
"context" | ||
|
||
"github.com/pkg/errors" | ||
"golang.org/x/oauth2" | ||
) | ||
|
||
type Identity struct { | ||
K8sNamespace string | ||
K8sServiceAccount string | ||
} | ||
|
||
type TokenSourceFactory interface { | ||
GetTokenSource(ctx context.Context, identity Identity) (oauth2.TokenSource, error) | ||
} | ||
|
||
func NewTokenSourceFactory(config TokenSourceFactoryConfig) (TokenSourceFactory, error) { | ||
if config.Type == TokenSourceTypeDefault { | ||
return NewDefaultTokenSourceFactory() | ||
} | ||
|
||
return nil, errors.Errorf("unknown token source type [%v], possible values are: 'default'", config.Type) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,74 @@ | ||
// Package bigquery implements WebAPI plugin for Google BigQuery | ||
package bigquery | ||
|
||
import ( | ||
"time" | ||
|
||
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/google" | ||
|
||
pluginsConfig "github.com/flyteorg/flyteplugins/go/tasks/config" | ||
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core" | ||
"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/webapi" | ||
"github.com/flyteorg/flytestdlib/config" | ||
) | ||
|
||
//go:generate pflags Config --default-var=defaultConfig | ||
|
||
var ( | ||
defaultConfig = Config{ | ||
WebAPI: webapi.PluginConfig{ | ||
ResourceQuotas: map[core.ResourceNamespace]int{ | ||
"default": 1000, | ||
}, | ||
ReadRateLimiter: webapi.RateLimiterConfig{ | ||
Burst: 100, | ||
QPS: 10, | ||
}, | ||
WriteRateLimiter: webapi.RateLimiterConfig{ | ||
Burst: 100, | ||
QPS: 10, | ||
}, | ||
Caching: webapi.CachingConfig{ | ||
Size: 500000, | ||
ResyncInterval: config.Duration{Duration: 30 * time.Second}, | ||
Workers: 10, | ||
MaxSystemFailures: 5, | ||
}, | ||
ResourceMeta: nil, | ||
}, | ||
ResourceConstraints: core.ResourceConstraintsSpec{ | ||
ProjectScopeResourceConstraint: &core.ResourceConstraint{ | ||
Value: 100, | ||
}, | ||
NamespaceScopeResourceConstraint: &core.ResourceConstraint{ | ||
Value: 50, | ||
}, | ||
}, | ||
GoogleTokenSource: google.GetDefaultConfig(), | ||
} | ||
|
||
configSection = pluginsConfig.MustRegisterSubSection("bigquery", &defaultConfig) | ||
) | ||
|
||
// Config is config for 'bigquery' plugin | ||
type Config struct { | ||
// WebAPI defines config for the base WebAPI plugin | ||
WebAPI webapi.PluginConfig `json:"webApi" pflag:",Defines config for the base WebAPI plugin."` | ||
|
||
// ResourceConstraints defines resource constraints on how many executions to be created per project/overall at any given time | ||
ResourceConstraints core.ResourceConstraintsSpec `json:"resourceConstraints" pflag:"-,Defines resource constraints on how many executions to be created per project/overall at any given time."` | ||
|
||
// GoogleTokenSource configures token source for BigQuery client | ||
GoogleTokenSource google.TokenSourceFactoryConfig `json:"googleTokenSource" pflag:",Defines Google token source"` | ||
|
||
// bigQueryEndpoint overrides BigQuery client endpoint, only for testing | ||
bigQueryEndpoint string | ||
} | ||
|
||
func GetConfig() *Config { | ||
return configSection.GetConfig().(*Config) | ||
} | ||
|
||
func SetConfig(cfg *Config) error { | ||
return configSection.SetConfig(cfg) | ||
} |
55 changes: 55 additions & 0 deletions
55
flyteplugins/go/tasks/plugins/webapi/bigquery/config_flags.go
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
Oops, something went wrong.