kebandClaude Opus 5.5
6fab41703e
Release / Build and publish (push) Successful in 22s
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
294 lines
9.4 KiB
Go
294 lines
9.4 KiB
Go
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/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")
|
|
}
|
|
}
|