-
Notifications
You must be signed in to change notification settings - Fork 1.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Refactor database engines registration (#10074)
- Loading branch information
Showing
6 changed files
with
240 additions
and
78 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,108 @@ | ||
/* | ||
Copyright 2022 Gravitational, Inc. | ||
Licensed under the Apache License, Version 2.0 (the "License"); | ||
you may not use this file except in compliance with the License. | ||
You may obtain a copy of the License at | ||
http://www.apache.org/licenses/LICENSE-2.0 | ||
Unless required by applicable law or agreed to in writing, software | ||
distributed under the License is distributed on an "AS IS" BASIS, | ||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
See the License for the specific language governing permissions and | ||
limitations under the License. | ||
*/ | ||
|
||
package common | ||
|
||
import ( | ||
"context" | ||
"sync" | ||
|
||
"github.com/gravitational/teleport/lib/auth" | ||
"github.com/gravitational/trace" | ||
|
||
"github.com/jonboulle/clockwork" | ||
"github.com/sirupsen/logrus" | ||
) | ||
|
||
var ( | ||
// engines is a global database engines registry. | ||
engines map[string]EngineFn | ||
// enginesMu protects access to the global engines registry map. | ||
enginesMu sync.RWMutex | ||
) | ||
|
||
// EngineFn defines a database engine constructor function. | ||
type EngineFn func(EngineConfig) Engine | ||
|
||
// RegisterEngine registers a new engine constructor. | ||
func RegisterEngine(fn EngineFn, names ...string) { | ||
enginesMu.Lock() | ||
defer enginesMu.Unlock() | ||
if engines == nil { | ||
engines = make(map[string]EngineFn) | ||
} | ||
for _, name := range names { | ||
engines[name] = fn | ||
} | ||
} | ||
|
||
// GetEngine returns a new engine for the provided configuration. | ||
func GetEngine(name string, conf EngineConfig) (Engine, error) { | ||
if err := conf.CheckAndSetDefaults(); err != nil { | ||
return nil, trace.Wrap(err) | ||
} | ||
enginesMu.RLock() | ||
engineFn := engines[name] | ||
enginesMu.RUnlock() | ||
if engineFn == nil { | ||
return nil, trace.NotFound("database engine %q is not registered", name) | ||
} | ||
return engineFn(conf), nil | ||
} | ||
|
||
// EngineConfig is the common configuration every database engine uses. | ||
type EngineConfig struct { | ||
// Auth handles database access authentication. | ||
Auth Auth | ||
// Audit emits database access audit events. | ||
Audit Audit | ||
// AuthClient is the cluster auth server client. | ||
AuthClient *auth.Client | ||
// CloudClients provides access to cloud API clients. | ||
CloudClients CloudClients | ||
// Context is the database server close context. | ||
Context context.Context | ||
// Clock is the clock interface. | ||
Clock clockwork.Clock | ||
// Log is used for logging. | ||
Log logrus.FieldLogger | ||
} | ||
|
||
// CheckAndSetDefaults validates the config and sets default values. | ||
func (c *EngineConfig) CheckAndSetDefaults() error { | ||
if c.Auth == nil { | ||
return trace.BadParameter("engine config Auth is missing") | ||
} | ||
if c.Audit == nil { | ||
return trace.BadParameter("engine config Audit is missing") | ||
} | ||
if c.AuthClient == nil { | ||
return trace.BadParameter("engine config AuthClient is missing") | ||
} | ||
if c.CloudClients == nil { | ||
return trace.BadParameter("engine config CloudClients are missing") | ||
} | ||
if c.Context == nil { | ||
c.Context = context.Background() | ||
} | ||
if c.Clock == nil { | ||
c.Clock = clockwork.NewRealClock() | ||
} | ||
if c.Log == nil { | ||
c.Log = logrus.StandardLogger() | ||
} | ||
return nil | ||
} |
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,75 @@ | ||
/* | ||
Copyright 2022 Gravitational, Inc. | ||
Licensed under the Apache License, Version 2.0 (the "License"); | ||
you may not use this file except in compliance with the License. | ||
You may obtain a copy of the License at | ||
http://www.apache.org/licenses/LICENSE-2.0 | ||
Unless required by applicable law or agreed to in writing, software | ||
distributed under the License is distributed on an "AS IS" BASIS, | ||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
See the License for the specific language governing permissions and | ||
limitations under the License. | ||
*/ | ||
|
||
package common | ||
|
||
import ( | ||
"context" | ||
"testing" | ||
|
||
"github.com/gravitational/teleport/lib/auth" | ||
|
||
"github.com/gravitational/trace" | ||
"github.com/jonboulle/clockwork" | ||
"github.com/sirupsen/logrus" | ||
"github.com/stretchr/testify/require" | ||
) | ||
|
||
// TestRegisterEngine verifies database engine registration. | ||
func TestRegisterEngine(t *testing.T) { | ||
ec := EngineConfig{ | ||
Context: context.Background(), | ||
Clock: clockwork.NewFakeClock(), | ||
Log: logrus.StandardLogger(), | ||
Auth: &testAuth{}, | ||
Audit: &testAudit{}, | ||
AuthClient: &auth.Client{}, | ||
CloudClients: NewCloudClients(), | ||
} | ||
|
||
// No engine is registered initially. | ||
engine, err := GetEngine("test", ec) | ||
require.Nil(t, engine) | ||
require.IsType(t, trace.NotFound(""), err) | ||
|
||
// Register a "test" engine. | ||
RegisterEngine(func(ec EngineConfig) Engine { | ||
return &testEngine{ec: ec} | ||
}, "test") | ||
|
||
// Create the registered engine instance. | ||
engine, err = GetEngine("test", ec) | ||
require.NoError(t, err) | ||
require.NotNil(t, engine) | ||
|
||
// Verify it's the one we registered. | ||
engineInst, ok := engine.(*testEngine) | ||
require.True(t, ok) | ||
require.Equal(t, ec, engineInst.ec) | ||
} | ||
|
||
type testEngine struct { | ||
Engine | ||
ec EngineConfig | ||
} | ||
|
||
type testAudit struct { | ||
Audit | ||
} | ||
|
||
type testAuth struct { | ||
Auth | ||
} |
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