Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
commit
45fd80b418
130 files changed
+26739
No files matched your search
@@ -0,0 +1,319 @@
|
||||
package databricks
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
dbconfig "github.com/databricks/databricks-sdk-go/config"
|
||||
"github.com/databricks/databricks-sdk-go/service/workspace"
|
||||
|
||||
"git.kebcloud.no/keb/lazysecrets/internal/provider"
|
||||
)
|
||||
|
||||
func TestReadProfiles(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "databrickscfg")
|
||||
cfg := `[DEFAULT]
|
||||
host = https://adb-1.1.azuredatabricks.net
|
||||
auth_type = azure-cli
|
||||
|
||||
[INFRA-DEV]
|
||||
host = adb-2.2.azuredatabricks.net
|
||||
auth_type = azure-cli
|
||||
|
||||
[ACCOUNT]
|
||||
host = https://accounts.azuredatabricks.net
|
||||
auth_type = azure-cli
|
||||
|
||||
[acct-by-id]
|
||||
host = https://example.cloud.databricks.com
|
||||
account_id = 123
|
||||
|
||||
[oauth]
|
||||
host = https://adb-3.3.azuredatabricks.net
|
||||
auth_type = databricks-cli
|
||||
|
||||
[no-host]
|
||||
token = x
|
||||
`
|
||||
if err := os.WriteFile(path, []byte(cfg), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := readProfiles(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := []profile{
|
||||
{name: "DEFAULT", host: "https://adb-1.1.azuredatabricks.net", authType: "azure-cli"},
|
||||
{name: "INFRA-DEV", host: "adb-2.2.azuredatabricks.net", authType: "azure-cli"},
|
||||
{name: "oauth", host: "https://adb-3.3.azuredatabricks.net", authType: "databricks-cli"},
|
||||
}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("profiles = %+v\nwant %+v", got, want)
|
||||
}
|
||||
|
||||
if got, err := readProfiles(filepath.Join(t.TempDir(), "missing")); err != nil || got != nil {
|
||||
t.Errorf("missing file: %v, %v", got, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginCommand(t *testing.T) {
|
||||
for authType, want := range map[string][]string{
|
||||
"azure-cli": {"az", "login"},
|
||||
"databricks-cli": {"databricks", "auth", "login", "--profile", "p"},
|
||||
"": {"databricks", "auth", "login", "--profile", "p"},
|
||||
"pat": nil,
|
||||
} {
|
||||
if got := loginCommand(profile{name: "p", authType: authType}); !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("%q: %v, want %v", authType, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fakeKeyVault stands in for the Azure provider.
|
||||
type fakeKeyVault struct {
|
||||
secrets map[string]string // name → value, for vault URL kvURL
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
const kvURL = "https://kv-app-prod.vault.azure.net/"
|
||||
|
||||
func (k *fakeKeyVault) Objects(_ context.Context, c provider.Container, _ provider.ObjectType) ([]provider.Object, error) {
|
||||
k.calls++
|
||||
if k.err != nil {
|
||||
return nil, k.err
|
||||
}
|
||||
if c.Detail != kvURL {
|
||||
return nil, errors.New("wrong vault " + c.Detail)
|
||||
}
|
||||
var out []provider.Object
|
||||
for name := range k.secrets {
|
||||
out = append(out, provider.Object{Name: name})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (k *fakeKeyVault) Value(_ context.Context, o provider.Object) (provider.SecretValue, error) {
|
||||
k.calls++
|
||||
if k.err != nil {
|
||||
return provider.SecretValue{}, k.err
|
||||
}
|
||||
if o.Container.Detail != kvURL {
|
||||
return provider.SecretValue{}, errors.New("wrong vault " + o.Container.Detail)
|
||||
}
|
||||
return provider.NewSecretValue([]byte(k.secrets[o.Name])), nil
|
||||
}
|
||||
|
||||
// workspaceAPI serves the secrets API. The "akv" scope behaves like an Azure
|
||||
// Key Vault-backed scope whose secrets Databricks refuses to read.
|
||||
func workspaceAPI(t *testing.T, status int) *httptest.Server {
|
||||
t.Helper()
|
||||
write := func(w http.ResponseWriter, code int, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(code)
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if status != 0 {
|
||||
write(w, status, map[string]any{"error_code": "UNAUTHENTICATED", "message": "Token is expired"})
|
||||
return
|
||||
}
|
||||
q := r.URL.Query()
|
||||
switch {
|
||||
case r.URL.Path == "/api/2.0/secrets/scopes/list":
|
||||
write(w, 200, map[string]any{"scopes": []any{
|
||||
map[string]any{"name": "app", "backend_type": "DATABRICKS"},
|
||||
map[string]any{"name": "akv", "backend_type": "AZURE_KEYVAULT",
|
||||
"keyvault_metadata": map[string]any{"dns_name": kvURL, "resource_id": "/x/kv-app-prod"}},
|
||||
}})
|
||||
case r.URL.Path == "/api/2.0/secrets/list" && q.Get("scope") == "app":
|
||||
write(w, 200, map[string]any{"secrets": []any{
|
||||
map[string]any{"key": "storage-key", "last_updated_timestamp": 1767225600000},
|
||||
map[string]any{"key": "client-secret"},
|
||||
}})
|
||||
case r.URL.Path == "/api/2.0/secrets/list" && q.Get("scope") == "locked":
|
||||
write(w, 403, map[string]any{"error_code": "PERMISSION_DENIED", "message": "User does not have READ permission on scope locked"})
|
||||
case r.URL.Path == "/api/2.0/secrets/get" && q.Get("scope") == "app":
|
||||
write(w, 200, map[string]any{"key": q.Get("key"),
|
||||
"value": base64.StdEncoding.EncodeToString([]byte("value-of-" + q.Get("key")))})
|
||||
case strings.HasPrefix(r.URL.Path, "/api/2.0/secrets/") && q.Get("scope") == "akv":
|
||||
write(w, 400, map[string]any{"error_code": "BAD_REQUEST",
|
||||
"message": "Reading secrets from Azure KeyVault-backed scopes is not supported"})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func newTestProvider(t *testing.T, srv *httptest.Server, kv KeyVault) *Provider {
|
||||
t.Helper()
|
||||
return &Provider{
|
||||
profiles: []profile{{name: "prod", host: srv.URL, authType: "azure-cli"}},
|
||||
keyVault: kv,
|
||||
newClient: func(profile) (*workspace.SecretsAPI, error) {
|
||||
return secretsClient(&dbconfig.Config{Host: srv.URL, Token: "token"})
|
||||
},
|
||||
clients: map[string]*workspace.SecretsAPI{},
|
||||
vaults: map[string]string{},
|
||||
}
|
||||
}
|
||||
|
||||
func scopes(t *testing.T, p *Provider) map[string]provider.Container {
|
||||
t.Helper()
|
||||
srcs, _ := p.Sources(context.Background())
|
||||
cs, err := p.Containers(context.Background(), srcs[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out := map[string]provider.Container{}
|
||||
for _, c := range cs {
|
||||
out[c.Name] = c
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestScopesSecretsAndValues(t *testing.T) {
|
||||
p := newTestProvider(t, workspaceAPI(t, 0), nil)
|
||||
sc := scopes(t, p)
|
||||
if sc["app"].Detail != "Databricks-backed" || sc["akv"].Detail != "Azure Key Vault-backed: "+kvURL {
|
||||
t.Errorf("scope details = %q / %q", sc["app"].Detail, sc["akv"].Detail)
|
||||
}
|
||||
|
||||
objs, err := p.Objects(context.Background(), sc["app"], provider.TypeSecret)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(objs) != 2 || objs[0].Name != "client-secret" || objs[1].Name != "storage-key" {
|
||||
t.Fatalf("objects = %+v", objs)
|
||||
}
|
||||
if objs[1].Updated.IsZero() || objs[1].Reference != "{{secrets/app/storage-key}}" {
|
||||
t.Errorf("metadata = %+v", objs[1])
|
||||
}
|
||||
|
||||
v, err := p.Value(context.Background(), objs[1])
|
||||
if err != nil || v.Reveal() != "value-of-storage-key" {
|
||||
t.Errorf("value = %q, %v", v.Reveal(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAKVScopeFallsBackToKeyVault(t *testing.T) {
|
||||
kv := &fakeKeyVault{secrets: map[string]string{"db-password": "from-the-vault"}}
|
||||
p := newTestProvider(t, workspaceAPI(t, 0), kv)
|
||||
akv := scopes(t, p)["akv"]
|
||||
|
||||
objs, err := p.Objects(context.Background(), akv, provider.TypeSecret)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(objs) != 1 || objs[0].Name != "db-password" || objs[0].Reference != "{{secrets/akv/db-password}}" {
|
||||
t.Fatalf("objects = %+v", objs)
|
||||
}
|
||||
v, err := p.Value(context.Background(), objs[0])
|
||||
if err != nil || v.Reveal() != "from-the-vault" {
|
||||
t.Errorf("value = %q, %v", v.Reveal(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAKVFallbackForRememberedScope(t *testing.T) {
|
||||
// A favourite from an earlier session: scopes not listed yet.
|
||||
kv := &fakeKeyVault{secrets: map[string]string{"db-password": "from-the-vault"}}
|
||||
p := newTestProvider(t, workspaceAPI(t, 0), kv)
|
||||
akv := provider.Container{
|
||||
Source: provider.Source{Kind: provider.KindDatabricks, ID: "prod"},
|
||||
ID: "prod/akv", Name: "akv", Detail: "Azure Key Vault-backed: " + kvURL,
|
||||
}
|
||||
v, err := p.Value(context.Background(), provider.Object{Container: akv, Name: "db-password"})
|
||||
if err != nil || v.Reveal() != "from-the-vault" {
|
||||
t.Errorf("value = %q, %v", v.Reveal(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAKVFallbackFailureReportsBoth(t *testing.T) {
|
||||
kvAuth := &provider.AuthError{Kind: provider.KindAzure, Login: []string{"az", "login"}, Err: errors.New("expired")}
|
||||
p := newTestProvider(t, workspaceAPI(t, 0), &fakeKeyVault{err: kvAuth})
|
||||
akv := scopes(t, p)["akv"]
|
||||
|
||||
_, err := p.Value(context.Background(), provider.Object{Container: akv, Name: "db-password"})
|
||||
if err == nil || !strings.Contains(err.Error(), "not supported") || !strings.Contains(err.Error(), "Key Vault fallback") {
|
||||
t.Fatalf("err = %v", err)
|
||||
}
|
||||
var auth *provider.AuthError
|
||||
if !errors.As(err, &auth) {
|
||||
t.Error("Key Vault auth error not reachable for the sign-in prompt")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoFallbackForDatabricksBackedScopes(t *testing.T) {
|
||||
kv := &fakeKeyVault{}
|
||||
p := newTestProvider(t, workspaceAPI(t, 0), kv)
|
||||
app := scopes(t, p)["app"]
|
||||
if _, err := p.Value(context.Background(), provider.Object{Container: app, Name: "missing"}); err != nil {
|
||||
// app/get returns a value for any key; just make sure the vault wasn't asked
|
||||
t.Fatal(err)
|
||||
}
|
||||
if kv.calls != 0 {
|
||||
t.Errorf("Key Vault consulted for a Databricks-backed scope")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnauthorizedIsAuthError(t *testing.T) {
|
||||
p := newTestProvider(t, workspaceAPI(t, 401), nil)
|
||||
srcs, _ := p.Sources(context.Background())
|
||||
_, err := p.Containers(context.Background(), srcs[0])
|
||||
var auth *provider.AuthError
|
||||
if !errors.As(err, &auth) {
|
||||
t.Fatalf("err = %v, want AuthError", err)
|
||||
}
|
||||
if auth.Kind != provider.KindDatabricks || !reflect.DeepEqual(auth.Login, []string{"az", "login"}) {
|
||||
t.Errorf("auth = %+v", auth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassify(t *testing.T) {
|
||||
pr := profile{name: "p", authType: "databricks-cli"}
|
||||
var auth *provider.AuthError
|
||||
if !errors.As(classify(pr, errors.New("default auth: azure-cli: cannot get access token: run az login")), &auth) {
|
||||
t.Error("token failure not classified as auth")
|
||||
}
|
||||
if errors.As(classify(pr, errors.New("dial tcp: no such host")), &auth) {
|
||||
t.Error("network error classified as auth")
|
||||
}
|
||||
if errors.As(classify(pr, context.Canceled), &auth) {
|
||||
t.Error("cancellation classified as auth")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeniedSaysWhatToCheck(t *testing.T) {
|
||||
p := newTestProvider(t, workspaceAPI(t, 0), nil)
|
||||
_, err := p.Objects(context.Background(), provider.Container{Source: provider.Source{ID: "prod"}, Name: "locked"}, provider.TypeSecret)
|
||||
var d *provider.Denied
|
||||
// The test profile signs in through az: that's the first check.
|
||||
if !errors.As(err, &d) || !strings.Contains(err.Error(), "READ permission on scope locked\n\nCheck:\n• az is signed in as the account with access: profile prod uses your az login") ||
|
||||
!strings.Contains(err.Error(), "\n• you have access to the workspace") {
|
||||
t.Errorf("err %q", err)
|
||||
}
|
||||
if got := deniedChecks(profile{authType: "pat"}, "denied"); strings.Contains(got[0], "az") {
|
||||
t.Errorf("a token profile's checks %q", got)
|
||||
}
|
||||
if got := deniedChecks(profile{}, "Source IP address: 1.2.3.4 is blocked by Databricks IP ACL for workspace: 1"); !strings.Contains(got[0], "network") {
|
||||
t.Errorf("IP ACL checks %q", got)
|
||||
}
|
||||
|
||||
// Both refused: one message, then both lists of checks.
|
||||
err = combine(&provider.Denied{Err: errors.New("dbx"), Checks: []string{"a", "b"}},
|
||||
&provider.Denied{Err: errors.New("kv"), Checks: []string{"b", "c"}})
|
||||
if got := err.Error(); got != "dbx; Key Vault fallback: kv\n\nCheck:\n• a\n• b\n• c" {
|
||||
t.Errorf("combined %q", got)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user