diff --git a/auth/credentials.go b/auth/credentials.go index c123ddd8..f1cb97df 100644 --- a/auth/credentials.go +++ b/auth/credentials.go @@ -15,9 +15,11 @@ package auth import ( + "context" "encoding/json" "errors" "fmt" + "io" "io/fs" "net/http" "net/url" @@ -75,46 +77,52 @@ type Credential struct { Audience string `toml:"audience,omitempty"` } -func (t *Credential) Refresh() bool { +func (t *Credential) Refresh() error { switch t.Type { case TypeApiKey: - rsp, err := http.DefaultClient.Do(&http.Request{ - Method: http.MethodGet, - URL: (*url.URL)(&t.AuthURI), - Header: http.Header{ - "authorization": []string{"Bearer " + t.ApiKey}, - }, - }) - if err != nil || rsp.StatusCode != http.StatusOK { - return false + req, err := http.NewRequestWithContext(context.Background(), + http.MethodGet, (*url.URL)(&t.AuthURI).String(), nil) + if err != nil { + return fmt.Errorf("apikey refresh: %w", err) + } + req.Header.Set("Authorization", "Bearer "+t.ApiKey) + + rsp, err := http.DefaultClient.Do(req) + if err != nil { + return fmt.Errorf("apikey refresh: %w", err) } defer rsp.Body.Close() + if rsp.StatusCode != http.StatusOK { + return fmt.Errorf("apikey refresh: status %s", rsp.Status) + } var tokenResp struct { Token string `json:"access_token"` } if err := json.NewDecoder(rsp.Body).Decode(&tokenResp); err != nil { - return false + return err } t.Token = tokenResp.Token - return true + return nil case TypeToken: if err := refreshOauth(t); err != nil { - return false + return fmt.Errorf("oauth refresh: %w", err) } - return true + return nil } - return false + return fmt.Errorf("unsupported credential type: %s", t.Type) } +// GetAuthToken returns the current token, refreshing if needed. +// Must not be called while credMu is held (Refresh may acquire it via UpdateCreds). func (t *Credential) GetAuthToken() string { if t.Token != "" { return t.Token } - if t.Refresh() { + if err := t.Refresh(); err == nil { _ = UpdateCreds() return t.Token } @@ -127,18 +135,30 @@ var ( credentialErr error loaded sync.Once credPath string + credPathMu sync.Mutex + credMu sync.RWMutex ) -func init() { - var err error - credPath, err = internal.GetCredentialPath() - if err != nil { - panic(fmt.Sprintf("failed to get credential path: %s", err)) +func getCredPath() (string, error) { + credPathMu.Lock() + defer credPathMu.Unlock() + if credPath == "" { + var err error + credPath, err = internal.GetCredentialPath() + if err != nil { + return "", fmt.Errorf("failed to get credential path: %w", err) + } } + return credPath, nil } func loadCreds() ([]Credential, error) { - credFile, err := os.Open(credPath) + cp, err := getCredPath() + if err != nil { + return nil, fmt.Errorf("failed to get credential path: %w", err) + } + + credFile, err := os.Open(cp) if err != nil { if errors.Is(err, os.ErrNotExist) { return []Credential{}, nil @@ -163,6 +183,8 @@ func GetCredentials(u *url.URL) (*Credential, error) { return nil, err } + credMu.RLock() + defer credMu.RUnlock() for i, cred := range loadedCredentials { if cred.RegistryURL.Host == u.Host { return &loadedCredentials[i], nil @@ -174,7 +196,10 @@ func GetCredentials(u *url.URL) (*Credential, error) { func LoadCredentials() error { loaded.Do(func() { - loadedCredentials, credentialErr = loadCreds() + creds, err := loadCreds() + credMu.Lock() + loadedCredentials, credentialErr = creds, err + credMu.Unlock() }) return credentialErr } @@ -184,6 +209,9 @@ func AddCredential(cred Credential, allowOverwrite bool) error { return err } + credMu.Lock() + defer credMu.Unlock() + idx := slices.IndexFunc(loadedCredentials, func(c Credential) bool { return c.RegistryURL.Host == cred.RegistryURL.Host }) @@ -196,7 +224,7 @@ func AddCredential(cred Credential, allowOverwrite bool) error { } else { loadedCredentials = append(loadedCredentials, cred) } - return UpdateCreds() + return writeCreds() } func RemoveCredential(host Uri) error { @@ -204,6 +232,9 @@ func RemoveCredential(host Uri) error { return err } + credMu.Lock() + defer credMu.Unlock() + idx := slices.IndexFunc(loadedCredentials, func(c Credential) bool { return c.RegistryURL.Host == host.Host }) @@ -213,20 +244,22 @@ func RemoveCredential(host Uri) error { } loadedCredentials = append(loadedCredentials[:idx], loadedCredentials[idx+1:]...) - return UpdateCreds() + return writeCreds() } -func UpdateCreds() error { - if err := LoadCredentials(); err != nil { - return err +// writeCreds persists loadedCredentials to disk. +// Caller must hold credMu. +func writeCreds() error { + cp, err := getCredPath() + if err != nil { + return fmt.Errorf("failed to get credential path: %w", err) } - err := os.MkdirAll(filepath.Dir(credPath), 0o700) - if err != nil { + if err := os.MkdirAll(filepath.Dir(cp), 0o700); err != nil { return err } - f, err := os.OpenFile(credPath, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0o600) + f, err := os.OpenFile(cp, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0o600) if err != nil { return err } @@ -239,13 +272,28 @@ func UpdateCreds() error { }) } +func UpdateCreds() error { + if err := LoadCredentials(); err != nil { + return err + } + + credMu.Lock() + defer credMu.Unlock() + return writeCreds() +} + func PurgeCredentials() error { + cp, err := getCredPath() + if err != nil { + return fmt.Errorf("failed to get credential path: %w", err) + } + var fileList = []string{ "credentials.toml", "columnar.lic", } - prefix := filepath.Dir(credPath) + prefix := filepath.Dir(cp) for _, file := range fileList { fullPath := filepath.Join(prefix, file) @@ -268,8 +316,12 @@ var ( ErrLicenseAlreadyExists = errors.New("license already exists (use --force to overwrite)") ) -func LicensePath() string { - return filepath.Join(filepath.Dir(credPath), "columnar.lic") +func LicensePath() (string, error) { + cp, err := getCredPath() + if err != nil { + return "", err + } + return filepath.Join(filepath.Dir(cp), "columnar.lic"), nil } func InstallLicenseFromFile(srcPath string, force bool) error { @@ -277,7 +329,10 @@ func InstallLicenseFromFile(srcPath string, force bool) error { return ErrLicenseWrongFilename } - destPath := LicensePath() + destPath, err := LicensePath() + if err != nil { + return fmt.Errorf("failed to determine license path: %w", err) + } if !force { if _, err := os.Stat(destPath); err == nil { @@ -302,8 +357,13 @@ func InstallLicenseFromFile(srcPath string, force bool) error { } func FetchColumnarLicense(cred *Credential) error { - licensePath := filepath.Join(filepath.Dir(credPath), "columnar.lic") - _, err := os.Stat(licensePath) + cp, err := getCredPath() + if err != nil { + return fmt.Errorf("failed to get credential path: %w", err) + } + + licensePath := filepath.Join(filepath.Dir(cp), "columnar.lic") + _, err = os.Stat(licensePath) if err == nil { // license exists already return nil } @@ -332,7 +392,7 @@ func FetchColumnarLicense(cred *Credential) error { return fmt.Errorf("unsupported credential type: %s", cred.Type) } - req, err := http.NewRequest(http.MethodGet, licenseURI, nil) + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, licenseURI, nil) if err != nil { return err } @@ -355,14 +415,19 @@ func FetchColumnarLicense(cred *Credential) error { } } - licenseFile, err := os.OpenFile(licensePath, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0o600) + tmp, err := os.CreateTemp(filepath.Dir(licensePath), ".lic.*") if err != nil { return err } - defer licenseFile.Close() - if _, err = licenseFile.ReadFrom(resp.Body); err != nil { - licenseFile.Close() - os.Remove(licensePath) + tmpName := tmp.Name() + defer os.Remove(tmpName) + + if _, err := io.Copy(tmp, resp.Body); err != nil { + tmp.Close() + return fmt.Errorf("write license: %w", err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("close license temp file: %w", err) } - return err + return os.Rename(tmpName, licensePath) } diff --git a/auth/credentials_test.go b/auth/credentials_test.go index 68a01012..7bdf0f39 100644 --- a/auth/credentials_test.go +++ b/auth/credentials_test.go @@ -20,7 +20,6 @@ import ( "net/url" "os" "path/filepath" - "sync" "testing" "github.com/pelletier/go-toml/v2" @@ -166,8 +165,8 @@ func TestCredential_Refresh_ApiKey(t *testing.T) { ApiKey: "test-api-key", } - success := cred.Refresh() - assert.True(t, success) + err := cred.Refresh() + assert.NoError(t, err) assert.Equal(t, "new-token", cred.Token) }) @@ -184,8 +183,8 @@ func TestCredential_Refresh_ApiKey(t *testing.T) { ApiKey: "invalid-key", } - success := cred.Refresh() - assert.False(t, success) + err := cred.Refresh() + assert.Error(t, err) }) t.Run("failed refresh with apikey - invalid json", func(t *testing.T) { @@ -202,8 +201,8 @@ func TestCredential_Refresh_ApiKey(t *testing.T) { ApiKey: "test-api-key", } - success := cred.Refresh() - assert.False(t, success) + err := cred.Refresh() + assert.Error(t, err) }) } @@ -273,20 +272,10 @@ func TestLoadCreds(t *testing.T) { } func TestGetCredentials(t *testing.T) { - // Save original values and restore after test - origCredPath := credPath - origLoadedCredentials := loadedCredentials - defer func() { - credPath = origCredPath - loadedCredentials = origLoadedCredentials - }() - - // Reset loaded state - loaded = sync.Once{} + resetCredState(t) t.Run("get existing credentials", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) u, _ := url.Parse("https://example.com") testCreds := struct { @@ -314,11 +303,10 @@ func TestGetCredentials(t *testing.T) { }) // Reset for next test - loaded = sync.Once{} + ResetCredentialsForTesting() t.Run("return nil when credentials not found", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) // Create empty credentials file testCreds := struct { @@ -341,20 +329,10 @@ func TestGetCredentials(t *testing.T) { } func TestAddCredential(t *testing.T) { - // Save original values and restore after test - origCredPath := credPath - origLoadedCredentials := loadedCredentials - defer func() { - credPath = origCredPath - loadedCredentials = origLoadedCredentials - }() - - // Reset loaded state - loaded = sync.Once{} + resetCredState(t) t.Run("add new credential", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) u, _ := url.Parse("https://example.com") newCred := Credential{ @@ -374,12 +352,10 @@ func TestAddCredential(t *testing.T) { }) // Reset for next test - loaded = sync.Once{} - loadedCredentials = nil + ResetCredentialsForTesting() t.Run("return error when credential already exists", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) u, _ := url.Parse("https://example.com") cred := Credential{ @@ -399,12 +375,10 @@ func TestAddCredential(t *testing.T) { }) // Reset for next test - loaded = sync.Once{} - loadedCredentials = nil + ResetCredentialsForTesting() t.Run("overwrite existing credential when allowOverwrite is true", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) u, _ := url.Parse("https://example.com") originalCred := Credential{ @@ -441,20 +415,10 @@ func TestAddCredential(t *testing.T) { } func TestRemoveCredential(t *testing.T) { - // Save original values and restore after test - origCredPath := credPath - origLoadedCredentials := loadedCredentials - defer func() { - credPath = origCredPath - loadedCredentials = origLoadedCredentials - }() - - // Reset loaded state - loaded = sync.Once{} + resetCredState(t) t.Run("remove existing credential", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) u, _ := url.Parse("https://example.com") cred := Credential{ @@ -478,12 +442,10 @@ func TestRemoveCredential(t *testing.T) { }) // Reset for next test - loaded = sync.Once{} - loadedCredentials = nil + ResetCredentialsForTesting() t.Run("return error when credential not found", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) // Create empty credentials file testCreds := struct { @@ -506,20 +468,10 @@ func TestRemoveCredential(t *testing.T) { } func TestUpdateCreds(t *testing.T) { - // Save original values and restore after test - origCredPath := credPath - origLoadedCredentials := loadedCredentials - defer func() { - credPath = origCredPath - loadedCredentials = origLoadedCredentials - }() - - // Reset loaded state - loaded = sync.Once{} + resetCredState(t) t.Run("update credentials file", func(t *testing.T) { - tmpDir := t.TempDir() - credPath = filepath.Join(tmpDir, "credentials.toml") + credPath = withTempCredPath(t) u, _ := url.Parse("https://example.com") @@ -543,9 +495,8 @@ func TestUpdateCreds(t *testing.T) { _, err = os.Stat(credPath) assert.NoError(t, err) - // Reset sync.Once to read from file - loaded = sync.Once{} - loadedCredentials = nil + // Reset to read from file + ResetCredentialsForTesting() // Verify content creds, err := loadCreds() diff --git a/auth/oauth.go b/auth/oauth.go index b18ead03..945f20a8 100644 --- a/auth/oauth.go +++ b/auth/oauth.go @@ -15,6 +15,7 @@ package auth import ( + "context" "encoding/json" "fmt" "net/http" @@ -64,12 +65,12 @@ func fetch[T any](u *url.URL, dest *T) error { if err != nil { return err } + defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return fmt.Errorf("failed to fetch %s: %s", u.String(), resp.Status) } - defer resp.Body.Close() return json.NewDecoder(resp.Body).Decode(dest) } @@ -99,8 +100,11 @@ func refreshOauth(cred *Credential) error { } payload := values.Encode() - req, _ := http.NewRequest(http.MethodPost, cfg.TokenEndpoint.String(), + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, cfg.TokenEndpoint.String(), strings.NewReader(payload)) + if err != nil { + return fmt.Errorf("failed to build token request: %w", err) + } req.Header.Add("content-type", "application/x-www-form-urlencoded") resp, err := http.DefaultClient.Do(req) if err != nil { diff --git a/auth/oauth_test.go b/auth/oauth_test.go index 135e162f..114599c0 100644 --- a/auth/oauth_test.go +++ b/auth/oauth_test.go @@ -307,8 +307,8 @@ func TestCredential_Refresh_OAuth(t *testing.T) { RefreshToken: "test-refresh-token", } - success := cred.Refresh() - assert.True(t, success) + err := cred.Refresh() + assert.NoError(t, err) assert.Equal(t, "refreshed-token", cred.Token) }) @@ -321,8 +321,8 @@ func TestCredential_Refresh_OAuth(t *testing.T) { RefreshToken: "test-refresh-token", } - success := cred.Refresh() - assert.False(t, success) + err := cred.Refresh() + assert.Error(t, err) }) } diff --git a/auth/test_helpers.go b/auth/test_helpers.go index 13b541a9..6595d2e0 100644 --- a/auth/test_helpers.go +++ b/auth/test_helpers.go @@ -14,11 +14,12 @@ package auth -import "sync" +import ( + "path/filepath" + "sync" + "testing" +) -// Test helpers to manipulate internal state for testing - -// SetCredPathForTesting sets the credential path for testing purposes func SetCredPathForTesting(path string) (restore func()) { orig := credPath credPath = path @@ -27,14 +28,33 @@ func SetCredPathForTesting(path string) (restore func()) { } } -// ResetCredentialsForTesting resets the loaded credentials state for testing func ResetCredentialsForTesting() { loaded = sync.Once{} loadedCredentials = nil credentialErr = nil } -// GetCredPathForTesting returns the current credential path func GetCredPathForTesting() string { return credPath } + +func resetCredState(t *testing.T) { + origCredPath := credPath + origLoadedCredentials := loadedCredentials + t.Cleanup(func() { + credPath = origCredPath + loadedCredentials = origLoadedCredentials + loaded = sync.Once{} + credentialErr = nil + }) + loaded = sync.Once{} + loadedCredentials = nil + credentialErr = nil +} + +func withTempCredPath(t *testing.T) string { + orig := credPath + t.Cleanup(func() { credPath = orig }) + credPath = filepath.Join(t.TempDir(), "credentials.toml") + return credPath +} diff --git a/client.go b/client.go new file mode 100644 index 00000000..42f1f454 --- /dev/null +++ b/client.go @@ -0,0 +1,149 @@ +// Copyright 2026 Columnar Technologies 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 dbc + +import ( + "fmt" + "net/http" + "net/url" + "os" + "path/filepath" + "runtime" + "sync" + + "github.com/columnar-tech/dbc/auth" + "github.com/columnar-tech/dbc/internal" + "github.com/google/uuid" + machineid "github.com/zeroshade/machine-id" +) + +type clientConfig struct { + httpClient *http.Client + registries []Registry + userAgent string + baseURL string + credentialResolver func(*url.URL) (*auth.Credential, error) +} + +type Option func(*clientConfig) + +// Client is a dbc client for searching registries and managing ADBC drivers. +type Client struct { + httpClient *http.Client + registries []Registry + userAgent string + mid string + uid uuid.UUID + setupOnce sync.Once + credentialResolver func(*url.URL) (*auth.Credential, error) +} + +// NewClient creates a new driver registry client with the given options. +func NewClient(opts ...Option) (*Client, error) { + cfg := &clientConfig{ + registries: []Registry{ + {BaseURL: mustParseURL("https://dbc-cdn.columnar.tech")}, + {BaseURL: mustParseURL("https://" + auth.DefaultOauthURI())}, + }, + userAgent: fmt.Sprintf("dbc-cli/%s (%s; %s)", Version, runtime.GOOS, runtime.GOARCH), + } + + for _, opt := range opts { + opt(cfg) + } + + httpClient := cfg.httpClient + if httpClient == nil { + httpClient = &http.Client{ + Transport: &uaRoundTripper{ + RoundTripper: http.DefaultTransport, + userAgent: cfg.userAgent, + }, + } + } + + if cfg.baseURL != "" { + cfg.registries = []Registry{{BaseURL: mustParseURL(cfg.baseURL)}} + } + + credResolver := cfg.credentialResolver + if credResolver == nil { + credResolver = auth.GetCredentials + } + + return &Client{ + httpClient: httpClient, + registries: cfg.registries, + userAgent: cfg.userAgent, + credentialResolver: credResolver, + }, nil +} + +func (c *Client) setup() { + c.setupOnce.Do(func() { + c.mid, _ = machineid.ProtectedID() + + userdir, err := internal.GetUserConfigPath() + if err != nil { + c.uid = uuid.New() + return + } + + fp := filepath.Join(userdir, "uid.uuid") + data, err := os.ReadFile(fp) + if err == nil { + if err = c.uid.UnmarshalBinary(data); err == nil { + return + } + } + + c.uid = uuid.New() + if err = os.MkdirAll(filepath.Dir(fp), 0o700); err == nil { + if data, err = c.uid.MarshalBinary(); err == nil { + os.WriteFile(fp, data, 0o600) + } + } + }) +} + +func (c *Client) HTTPClient() *http.Client { return c.httpClient } + +// Registries returns the list of driver registries configured for this client. +func (c *Client) Registries() []Registry { return c.registries } + +// UserAgent returns the user agent string used by this client. +func (c *Client) UserAgent() string { return c.userAgent } + +// WithHTTPClient sets the HTTP client to use for requests. +func WithHTTPClient(hc *http.Client) Option { + return func(cfg *clientConfig) { cfg.httpClient = hc } +} + +// WithRegistries sets the driver registries to use. +func WithRegistries(r []Registry) Option { + return func(cfg *clientConfig) { cfg.registries = append([]Registry(nil), r...) } +} + +// WithBaseURL sets the base URL for the driver registry. +func WithBaseURL(u string) Option { + return func(cfg *clientConfig) { cfg.baseURL = u } +} + +// WithUserAgent sets the user agent string for requests. This only takes +// effect when no custom HTTP client is provided via WithHTTPClient; if a +// custom client is supplied its transport is used as-is. +func WithUserAgent(ua string) Option { + return func(cfg *clientConfig) { cfg.userAgent = ua } +} diff --git a/client_auth.go b/client_auth.go new file mode 100644 index 00000000..7734d9d1 --- /dev/null +++ b/client_auth.go @@ -0,0 +1,59 @@ +// Copyright 2026 Columnar Technologies 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 dbc + +import ( + "fmt" + "net/url" + + "github.com/columnar-tech/dbc/auth" +) + +// WithCredential sets a specific credential to use for all requests. +func WithCredential(cred *auth.Credential) Option { + credCopy := *cred + return func(cfg *clientConfig) { + cfg.credentialResolver = func(_ *url.URL) (*auth.Credential, error) { + return &credCopy, nil + } + } +} + +// WithAuthFromFilesystem configures the client to read credentials from the filesystem. +func WithAuthFromFilesystem() Option { + return func(cfg *clientConfig) { + cfg.credentialResolver = auth.GetCredentials + } +} + +func (c *Client) getCredential(u *url.URL) (*auth.Credential, error) { + if c.credentialResolver == nil { + return nil, nil + } + return c.credentialResolver(u) +} + +// Login saves a credential for the given registry. +func (c *Client) Login(cred *auth.Credential) error { + if cred == nil { + return fmt.Errorf("credential must not be nil") + } + return auth.AddCredential(*cred, true) +} + +// Logout removes the credential for the given registry URL. +func (c *Client) Logout(registryURL *url.URL) error { + return auth.RemoveCredential(auth.Uri(*registryURL)) +} diff --git a/client_auth_test.go b/client_auth_test.go new file mode 100644 index 00000000..ed06c736 --- /dev/null +++ b/client_auth_test.go @@ -0,0 +1,167 @@ +// Copyright 2026 Columnar Technologies 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 dbc_test + +import ( + "net/http" + "net/http/httptest" + "net/url" + "path/filepath" + "testing" + + "github.com/columnar-tech/dbc" + "github.com/columnar-tech/dbc/auth" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestClientWithCredential(t *testing.T) { + const token = "test-injected-token" + + var gotAuthHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuthHeader = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/yaml") + w.WriteHeader(http.StatusOK) + w.Write([]byte("drivers: []\n")) + })) + defer srv.Close() + + srvURL, err := url.Parse(srv.URL) + require.NoError(t, err) + + cred := &auth.Credential{ + Type: auth.TypeApiKey, + ApiKey: "unused-api-key", + Token: token, + RegistryURL: auth.Uri(*srvURL), + } + + c, err := dbc.NewClient( + dbc.WithHTTPClient(&http.Client{}), + dbc.WithBaseURL(srv.URL), + dbc.WithCredential(cred), + ) + require.NoError(t, err) + + _, _ = c.Search("") + + assert.Equal(t, "Bearer "+token, gotAuthHeader) +} + +func TestClientWithAuthFromFilesystem(t *testing.T) { + var gotAuthHeader string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuthHeader = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/yaml") + w.Write([]byte("drivers: []\n")) + })) + defer srv.Close() + + tmpDir := t.TempDir() + restore := auth.SetCredPathForTesting(filepath.Join(tmpDir, "credentials.toml")) + defer restore() + auth.ResetCredentialsForTesting() + + c, err := dbc.NewClient( + dbc.WithHTTPClient(&http.Client{}), + dbc.WithBaseURL(srv.URL), + dbc.WithAuthFromFilesystem(), + ) + require.NoError(t, err) + + _, _ = c.Search("") + + assert.Empty(t, gotAuthHeader) +} + +func TestClientLogin(t *testing.T) { + tmpDir := t.TempDir() + credPath := filepath.Join(tmpDir, "credentials.toml") + restore := auth.SetCredPathForTesting(credPath) + defer restore() + auth.ResetCredentialsForTesting() + + c, err := dbc.NewClient() + require.NoError(t, err) + + u, err := url.Parse("https://login-test.example.com") + require.NoError(t, err) + cred := &auth.Credential{ + Type: auth.TypeApiKey, + ApiKey: "my-api-key", + Token: "my-token", + RegistryURL: auth.Uri(*u), + } + + err = c.Login(cred) + require.NoError(t, err) + + auth.ResetCredentialsForTesting() + found, err := auth.GetCredentials(u) + require.NoError(t, err) + require.NotNil(t, found) + assert.Equal(t, "my-api-key", found.ApiKey) + + cred2 := &auth.Credential{ + Type: auth.TypeApiKey, + ApiKey: "updated-key", + Token: "updated-token", + RegistryURL: auth.Uri(*u), + } + err = c.Login(cred2) + require.NoError(t, err) + + auth.ResetCredentialsForTesting() + updated, err := auth.GetCredentials(u) + require.NoError(t, err) + require.NotNil(t, updated) + assert.Equal(t, "updated-key", updated.ApiKey) +} + +func TestClientLogout(t *testing.T) { + tmpDir := t.TempDir() + credPath := filepath.Join(tmpDir, "credentials.toml") + restore := auth.SetCredPathForTesting(credPath) + defer restore() + auth.ResetCredentialsForTesting() + + c, err := dbc.NewClient() + require.NoError(t, err) + + u, err := url.Parse("https://logout-test.example.com") + require.NoError(t, err) + cred := &auth.Credential{ + Type: auth.TypeApiKey, + ApiKey: "my-api-key", + Token: "my-token", + RegistryURL: auth.Uri(*u), + } + + require.NoError(t, c.Login(cred)) + + err = c.Logout(u) + require.NoError(t, err) + + auth.ResetCredentialsForTesting() + found, err := auth.GetCredentials(u) + require.NoError(t, err) + assert.Nil(t, found) + + u2, err := url.Parse("https://not-registered.example.com") + require.NoError(t, err) + err = c.Logout(u2) + assert.Error(t, err) +} diff --git a/client_config.go b/client_config.go new file mode 100644 index 00000000..9dc2845e --- /dev/null +++ b/client_config.go @@ -0,0 +1,37 @@ +// Copyright 2026 Columnar Technologies 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 dbc + +import "github.com/columnar-tech/dbc/config" + +// GetConfig returns the configuration at the specified level. +func (c *Client) GetConfig(level config.ConfigLevel) config.Config { + return config.Get()[level] +} + +// ListInstalled returns a list of installed drivers at the specified configuration level. +func (c *Client) ListInstalled(level config.ConfigLevel) []config.DriverInfo { + return config.FindDriverConfigs(level) +} + +// GetDriver retrieves driver information from the given configuration. +func (c *Client) GetDriver(cfg config.Config, name string) (config.DriverInfo, error) { + return config.GetDriver(cfg, name) +} + +// CreateManifest creates a manifest file for the given driver configuration. +func (c *Client) CreateManifest(cfg config.Config, di config.DriverInfo) error { + return config.CreateManifest(cfg, di) +} diff --git a/client_config_test.go b/client_config_test.go new file mode 100644 index 00000000..7ccc6b94 --- /dev/null +++ b/client_config_test.go @@ -0,0 +1,200 @@ +// Copyright 2026 Columnar Technologies 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. + +//go:build !windows + +package dbc_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/Masterminds/semver/v3" + "github.com/columnar-tech/dbc" + "github.com/columnar-tech/dbc/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const clientTestManifestTOML = ` +name = 'Test Driver' +publisher = 'Test Publisher' +license = 'MIT' +version = '1.2.3' +source = 'dbc' + +[ADBC] +version = '1.1.0' + +[Driver] +entrypoint = 'AdbcDriverInit' + +[Driver.shared] +linux_amd64 = '/path/to/driver.so' +` + +func newTestClient(t *testing.T) *dbc.Client { + t.Helper() + c, err := dbc.NewClient() + require.NoError(t, err) + return c +} + +func makeClientTestDriverInfo(id string, filePath string) config.DriverInfo { + di := config.DriverInfo{ + ID: id, + FilePath: filePath, + Name: "Test Driver", + Publisher: "Test Publisher", + License: "MIT", + Version: semver.MustParse("1.2.3"), + Source: "dbc", + } + di.Driver.Entrypoint = "AdbcDriverInit" + di.Driver.Shared.Set("linux_amd64", filepath.Join(filePath, id, "driver.so")) + return di +} + +func TestClientGetConfig(t *testing.T) { + t.Run("returns_config_for_env_level", func(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.WriteFile( + filepath.Join(tmpDir, "mydriver.toml"), + []byte(clientTestManifestTOML), + 0644, + )) + t.Setenv("ADBC_DRIVER_PATH", tmpDir) + + c := newTestClient(t) + cfg := c.GetConfig(config.ConfigEnv) + + assert.Equal(t, config.ConfigEnv, cfg.Level) + assert.True(t, cfg.Exists) + }) + + t.Run("returns_config_for_system_level", func(t *testing.T) { + c := newTestClient(t) + cfg := c.GetConfig(config.ConfigSystem) + + assert.Equal(t, config.ConfigSystem, cfg.Level) + }) +} + +func TestClientListInstalled(t *testing.T) { + t.Run("returns_installed_drivers", func(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("ADBC_DRIVER_PATH", tmpDir) + + require.NoError(t, os.WriteFile( + filepath.Join(tmpDir, "driver1.toml"), + []byte(clientTestManifestTOML), + 0644, + )) + require.NoError(t, os.WriteFile( + filepath.Join(tmpDir, "driver2.toml"), + []byte(clientTestManifestTOML), + 0644, + )) + + c := newTestClient(t) + drivers := c.ListInstalled(config.ConfigEnv) + require.Len(t, drivers, 2) + + ids := make([]string, len(drivers)) + for i, d := range drivers { + ids[i] = d.ID + } + assert.ElementsMatch(t, []string{"driver1", "driver2"}, ids) + }) + + t.Run("empty_path_returns_empty_slice", func(t *testing.T) { + t.Setenv("ADBC_DRIVER_PATH", "") + t.Setenv("VIRTUAL_ENV", "") + t.Setenv("CONDA_PREFIX", "") + + c := newTestClient(t) + drivers := c.ListInstalled(config.ConfigEnv) + assert.Empty(t, drivers) + }) +} + +func TestClientGetDriver(t *testing.T) { + t.Run("found_in_env_config", func(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.WriteFile( + filepath.Join(tmpDir, "mydriver.toml"), + []byte(clientTestManifestTOML), + 0644, + )) + + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + c := newTestClient(t) + di, err := c.GetDriver(cfg, "mydriver") + require.NoError(t, err) + assert.Equal(t, "mydriver", di.ID) + assert.Equal(t, "Test Driver", di.Name) + assert.Equal(t, "1.2.3", di.Version.String()) + }) + + t.Run("not_found_returns_error", func(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + c := newTestClient(t) + _, err := c.GetDriver(cfg, "nonexistent") + assert.Error(t, err) + }) +} + +func TestClientCreateManifest(t *testing.T) { + t.Run("creates_manifest_file", func(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + di := makeClientTestDriverInfo("mydriver", tmpDir) + + c := newTestClient(t) + err := c.CreateManifest(cfg, di) + require.NoError(t, err) + + assert.FileExists(t, filepath.Join(tmpDir, "mydriver.toml")) + }) + + t.Run("creates_location_if_absent", func(t *testing.T) { + newDir := filepath.Join(t.TempDir(), "nonexistent", "subdir") + cfg := config.Config{ + Level: config.ConfigEnv, + Location: newDir, + } + + di := makeClientTestDriverInfo("newdriver", newDir) + + c := newTestClient(t) + err := c.CreateManifest(cfg, di) + require.NoError(t, err) + + assert.FileExists(t, filepath.Join(newDir, "newdriver.toml")) + }) +} diff --git a/client_methods.go b/client_methods.go new file mode 100644 index 00000000..ce1640d8 --- /dev/null +++ b/client_methods.go @@ -0,0 +1,305 @@ +// Copyright 2026 Columnar Technologies 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 dbc + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "path" + "path/filepath" + "sort" + "strings" + + "github.com/columnar-tech/dbc/auth" + "github.com/columnar-tech/dbc/config" + "github.com/go-faster/yaml" +) + +func (c *Client) makeRequest(u string) (*http.Response, error) { + c.setup() + + uri, err := url.Parse(u) + if err != nil { + return nil, fmt.Errorf("failed to parse URL %s: %w", u, err) + } + + cred, err := c.getCredential(uri) + if err != nil && !os.IsNotExist(err) { + return nil, fmt.Errorf("failed to read credentials: %w", err) + } + + q := uri.Query() + q.Add("mid", c.mid) + q.Add("uid", c.uid.String()) + uri.RawQuery = q.Encode() + + buildReq := func(token string) (*http.Request, error) { + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, uri.String(), nil) + if err != nil { + return nil, err + } + if uri.Path == "/index.yaml" { + req.Header.Set("Accept", "application/yaml") + } + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + return req, nil + } + + token := "" + if cred != nil { + if auth.IsColumnarPrivateRegistry(uri) { + _ = auth.FetchColumnarLicense(cred) + } + token = cred.GetAuthToken() + } + + req, err := buildReq(token) + if err != nil { + return nil, fmt.Errorf("failed to build request: %w", err) + } + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + + if resp.StatusCode == http.StatusUnauthorized && cred != nil { + resp.Body.Close() + if err := cred.Refresh(); err != nil { + return nil, fmt.Errorf("failed to refresh auth token: %w", err) + } + req, err = buildReq(cred.GetAuthToken()) + if err != nil { + return nil, fmt.Errorf("failed to build retry request: %w", err) + } + resp, err = c.httpClient.Do(req) + if err != nil { + return nil, err + } + } + + switch resp.StatusCode { + case http.StatusUnauthorized, http.StatusForbidden: + err = ErrUnauthorized + if auth.IsColumnarPrivateRegistry(uri) && cred != nil { + err = ErrUnauthorizedColumnar + } + resp.Body.Close() + return nil, fmt.Errorf("%s%s: %w", uri.Host, uri.Path, err) + } + + return resp, nil +} + +func (c *Client) getDriverListFromIndex(index *Registry) ([]Driver, error) { + resp, err := c.makeRequest(index.BaseURL.JoinPath("/index.yaml").String()) + if err != nil { + return nil, fmt.Errorf("failed to fetch drivers: %w", err) + } + + if resp.StatusCode != http.StatusOK { + resp.Body.Close() + return nil, fmt.Errorf("failed to fetch drivers: %s", resp.Status) + } + + defer resp.Body.Close() + drivers := struct { + Name string `yaml:"name"` + Drivers []Driver `yaml:"drivers"` + }{} + + if err = yaml.NewDecoder(resp.Body).Decode(&drivers); err != nil { + return nil, fmt.Errorf("failed to parse driver registry index: %s", err) + } + + if drivers.Name != "" { + index.Name = drivers.Name + } + + for i := range drivers.Drivers { + drivers.Drivers[i].Registry = index + } + + result := drivers.Drivers + sort.Slice(result, func(i, j int) bool { + return result[i].Path < result[j].Path + }) + + return result, nil +} + +// Search searches for drivers matching the given pattern across all registries. +func (c *Client) Search(pattern string) ([]Driver, error) { + var ( + allDrivers []Driver + totalErr error + ) + + for i := range c.registries { + drivers, err := c.getDriverListFromIndex(&c.registries[i]) + if err != nil { + totalErr = errors.Join(totalErr, fmt.Errorf("registry %s: %w", c.registries[i].BaseURL, err)) + continue + } + c.registries[i].Drivers = drivers + allDrivers = append(allDrivers, drivers...) + } + + if pattern == "" { + return allDrivers, totalErr + } + + lowerPattern := strings.ToLower(pattern) + var filtered []Driver + for _, d := range allDrivers { + if strings.Contains(strings.ToLower(d.Path), lowerPattern) || + strings.Contains(strings.ToLower(d.Title), lowerPattern) { + filtered = append(filtered, d) + } + } + + return filtered, totalErr +} + +func (c *Client) downloadPackage(pkg PkgInfo) (*os.File, error) { + if pkg.Path == nil { + return nil, fmt.Errorf("cannot download package for %s: no url set", pkg.Driver.Title) + } + + location := pkg.Path.String() + rsp, err := c.makeRequest(location) + if err != nil { + return nil, fmt.Errorf("failed to download driver: %w", err) + } + + if rsp.StatusCode != http.StatusOK { + rsp.Body.Close() + return nil, fmt.Errorf("failed to download driver %s: %s", location, rsp.Status) + } + defer rsp.Body.Close() + + fname := path.Base(location) + tmpdir, err := os.MkdirTemp(os.TempDir(), "adbc-drivers-*") + if err != nil { + return nil, fmt.Errorf("failed to create temp dir: %w", err) + } + + var output *os.File + defer func() { + if output == nil { + os.RemoveAll(tmpdir) + } + }() + + output, err = os.Create(path.Join(tmpdir, fname)) + if err != nil { + return nil, fmt.Errorf("failed to create temp file to download to: %w", err) + } + + if _, err = io.Copy(output, rsp.Body); err != nil { + output.Close() + output = nil + return nil, fmt.Errorf("failed to write driver file: %w", err) + } + + return output, nil +} + +// Download fetches the tarball for pkg and returns its contents as an +// io.ReadCloser. The caller is responsible for closing the returned body. +// Auth credentials are resolved and injected automatically, including token +// refresh on 401. +func (c *Client) Download(pkg PkgInfo) (io.ReadCloser, error) { + if pkg.Path == nil { + return nil, fmt.Errorf("cannot download package for %s: no url set", pkg.Driver.Title) + } + rsp, err := c.makeRequest(pkg.Path.String()) + if err != nil { + return nil, fmt.Errorf("failed to download %s: %w", pkg.Path, err) + } + if rsp.StatusCode != http.StatusOK { + defer rsp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(rsp.Body, 1024)) + if len(body) > 0 { + return nil, fmt.Errorf("failed to download %s: %s: %s", pkg.Path, rsp.Status, body) + } + return nil, fmt.Errorf("failed to download %s: %s", pkg.Path, rsp.Status) + } + return rsp.Body, nil +} + +// Install installs a driver with the given name to the specified configuration. +func (c *Client) Install(cfg config.Config, driverName string) (*config.Manifest, error) { + drivers, err := c.Search(driverName) + // Only fail if the driver wasn't found in any registry; partial registry errors + // are acceptable as long as we can still locate the target driver. + if err != nil && len(drivers) == 0 { + return nil, fmt.Errorf("failed to search for driver %s: %w", driverName, err) + } + + var found *Driver + for i := range drivers { + if drivers[i].Path == driverName { + found = &drivers[i] + break + } + } + + if found == nil { + return nil, fmt.Errorf("driver %q not found", driverName) + } + + pkg, err := found.GetPackage(nil, config.PlatformTuple(), false) + if err != nil { + return nil, fmt.Errorf("failed to get package for driver %s: %w", driverName, err) + } + + f, err := c.downloadPackage(pkg) + if err != nil { + return nil, fmt.Errorf("failed to download driver %s: %w", driverName, err) + } + defer os.RemoveAll(filepath.Dir(f.Name())) + + manifest, err := config.InstallDriver(cfg, driverName, f) + if err != nil { + return nil, fmt.Errorf("failed to install driver %s: %w", driverName, err) + } + + if err := config.CreateManifest(cfg, manifest.DriverInfo); err != nil { + return nil, fmt.Errorf("failed to create manifest for driver %s: %w", driverName, err) + } + + return &manifest, nil +} + +// Uninstall uninstalls a driver with the given name from the specified configuration. +func (c *Client) Uninstall(cfg config.Config, driverName string) error { + di, err := config.GetDriver(cfg, driverName) + if err != nil { + return fmt.Errorf("failed to find driver %q: %w", driverName, err) + } + + if err := config.UninstallDriver(cfg, di); err != nil { + return fmt.Errorf("failed to uninstall driver %q: %w", driverName, err) + } + + return nil +} diff --git a/client_methods_test.go b/client_methods_test.go new file mode 100644 index 00000000..78c1afa2 --- /dev/null +++ b/client_methods_test.go @@ -0,0 +1,224 @@ +// Copyright 2026 Columnar Technologies 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 dbc_test + +import ( + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/columnar-tech/dbc" + "github.com/columnar-tech/dbc/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTestClientForServer(t *testing.T, serverURL string) *dbc.Client { + t.Helper() + c, err := dbc.NewClient( + dbc.WithHTTPClient(&http.Client{}), + dbc.WithBaseURL(serverURL), + ) + require.NoError(t, err) + return c +} + +func newInstallTestServer(t *testing.T) *httptest.Server { + t.Helper() + + indexData, err := os.ReadFile(filepath.Join("cmd", "dbc", "testdata", "test_index.yaml")) + require.NoError(t, err) + + tarballData, err := os.ReadFile(filepath.Join("cmd", "dbc", "testdata", "test-driver-1.tar.gz")) + require.NoError(t, err) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/index.yaml": + w.Header().Set("Content-Type", "application/yaml") + w.Write(indexData) + case strings.HasSuffix(r.URL.Path, ".tar.gz"): + w.Header().Set("Content-Type", "application/gzip") + w.Header().Set("Content-Length", fmt.Sprint(len(tarballData))) + w.Write(tarballData) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestClientSearch(t *testing.T) { + srv := newInstallTestServer(t) + c, err := dbc.NewClient( + dbc.WithHTTPClient(&http.Client{}), + dbc.WithBaseURL(srv.URL), + ) + require.NoError(t, err) + + t.Run("empty pattern returns all drivers", func(t *testing.T) { + drivers, err := c.Search("") + require.NoError(t, err) + assert.NotEmpty(t, drivers) + + var paths []string + for _, d := range drivers { + paths = append(paths, d.Path) + } + assert.Contains(t, paths, "test-driver-1") + assert.Contains(t, paths, "test-driver-2") + }) + + t.Run("pattern matches by path", func(t *testing.T) { + drivers, err := c.Search("test-driver-1") + require.NoError(t, err) + require.NotEmpty(t, drivers) + + var paths []string + for _, d := range drivers { + paths = append(paths, d.Path) + } + assert.Contains(t, paths, "test-driver-1") + assert.NotContains(t, paths, "test-driver-2") + }) + + t.Run("nonexistent pattern returns empty list", func(t *testing.T) { + drivers, err := c.Search("nonexistent-driver-xyz") + require.NoError(t, err) + assert.Empty(t, drivers) + }) +} + +func TestClientInstall(t *testing.T) { + srv := newInstallTestServer(t) + c := newTestClientForServer(t, srv.URL) + + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + t.Run("installs driver successfully", func(t *testing.T) { + manifest, err := c.Install(cfg, "test-driver-1") + require.NoError(t, err) + require.NotNil(t, manifest) + assert.Equal(t, "test-driver-1", manifest.DriverInfo.ID) + assert.NotNil(t, manifest.DriverInfo.Version) + }) + + t.Run("returns error for nonexistent driver", func(t *testing.T) { + _, err := c.Install(cfg, "nonexistent-driver") + assert.Error(t, err) + assert.Contains(t, err.Error(), "nonexistent-driver") + }) +} + +func TestClientUninstall(t *testing.T) { + srv := newInstallTestServer(t) + c := newTestClientForServer(t, srv.URL) + + t.Run("uninstalls a previously installed driver", func(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + _, err := c.Install(cfg, "test-driver-1") + require.NoError(t, err) + + manifestPath := filepath.Join(tmpDir, "test-driver-1.toml") + _, err = os.Stat(manifestPath) + require.NoError(t, err, "manifest TOML should exist after install") + + err = c.Uninstall(cfg, "test-driver-1") + require.NoError(t, err) + + _, err = os.Stat(manifestPath) + assert.True(t, os.IsNotExist(err), "manifest TOML should be removed after uninstall") + }) + + t.Run("returns error for driver not installed", func(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + err := c.Uninstall(cfg, "not-installed-driver") + assert.Error(t, err) + assert.Contains(t, err.Error(), "not-installed-driver") + }) +} + +func TestClientDownload(t *testing.T) { + srv := newInstallTestServer(t) + c := newTestClientForServer(t, srv.URL) + + t.Run("successful download returns readable body", func(t *testing.T) { + pkgURL, err := url.Parse(srv.URL + "/test-driver-1.tar.gz") + require.NoError(t, err) + + pkg := dbc.PkgInfo{ + Driver: dbc.Driver{Title: "test-driver-1"}, + Path: pkgURL, + } + body, err := c.Download(pkg) + require.NoError(t, err) + defer body.Close() + + data, err := io.ReadAll(body) + require.NoError(t, err) + assert.NotEmpty(t, data, "download body should not be empty") + }) + + t.Run("nil path returns error", func(t *testing.T) { + pkg := dbc.PkgInfo{ + Driver: dbc.Driver{Title: "nil-path-driver"}, + } + _, err := c.Download(pkg) + require.Error(t, err) + assert.Contains(t, err.Error(), "no url set") + assert.Contains(t, err.Error(), "nil-path-driver") + }) + + t.Run("non-200 status returns error with body", func(t *testing.T) { + errSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + fmt.Fprint(w, "package not found in storage backend") + })) + t.Cleanup(errSrv.Close) + + errClient := newTestClientForServer(t, errSrv.URL) + pkgURL, err := url.Parse(errSrv.URL + "/some-driver.tar.gz") + require.NoError(t, err) + + pkg := dbc.PkgInfo{ + Driver: dbc.Driver{Title: "error-driver"}, + Path: pkgURL, + } + _, err = errClient.Download(pkg) + require.Error(t, err) + assert.Contains(t, err.Error(), "500") + assert.Contains(t, err.Error(), "package not found in storage backend") + }) +} diff --git a/client_test.go b/client_test.go new file mode 100644 index 00000000..cae9c987 --- /dev/null +++ b/client_test.go @@ -0,0 +1,63 @@ +// Copyright 2026 Columnar Technologies 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 dbc_test + +import ( + "net/http" + "net/url" + "testing" + + "github.com/columnar-tech/dbc" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewClientDefaults(t *testing.T) { + c, err := dbc.NewClient() + require.NoError(t, err) + assert.NotNil(t, c) + assert.NotNil(t, c.HTTPClient()) + assert.Len(t, c.Registries(), 2) +} + +func TestNewClientWithHTTPClient(t *testing.T) { + custom := &http.Client{} + c, err := dbc.NewClient(dbc.WithHTTPClient(custom)) + require.NoError(t, err) + assert.Same(t, custom, c.HTTPClient()) +} + +func TestNewClientWithRegistries(t *testing.T) { + u, err := url.Parse("https://example.com") + require.NoError(t, err) + regs := []dbc.Registry{{BaseURL: u}} + c, err := dbc.NewClient(dbc.WithRegistries(regs)) + require.NoError(t, err) + assert.Equal(t, regs, c.Registries()) +} + +func TestNewClientWithBaseURL(t *testing.T) { + c, err := dbc.NewClient(dbc.WithBaseURL("https://custom.example.com")) + require.NoError(t, err) + regs := c.Registries() + require.Len(t, regs, 1) + assert.Equal(t, "https://custom.example.com", regs[0].BaseURL.String()) +} + +func TestNewClientWithUserAgent(t *testing.T) { + c, err := dbc.NewClient(dbc.WithUserAgent("custom-agent/1.0")) + require.NoError(t, err) + assert.Equal(t, "custom-agent/1.0", c.UserAgent()) +} diff --git a/cmd/dbc/add.go b/cmd/dbc/add.go index 48e184d7..f11392ca 100644 --- a/cmd/dbc/add.go +++ b/cmd/dbc/add.go @@ -15,7 +15,6 @@ package main import ( - "errors" "fmt" "os" "path/filepath" @@ -58,13 +57,10 @@ func (c AddCmd) GetModelCustom(baseModel baseModel) tea.Model { func (c AddCmd) GetModel() tea.Model { return addModel{ - Driver: c.Driver, - Path: c.Path, - Pre: c.Pre, - baseModel: baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, + Driver: c.Driver, + Path: c.Path, + Pre: c.Pre, + baseModel: defaultBaseModel(), } } @@ -109,20 +105,10 @@ func (m addModel) Init() tea.Cmd { return err } - f, err := os.Open(p) + m.list, err = openAndDecodeDriverList(m.Path) if err != nil { - if errors.Is(err, os.ErrNotExist) { - return fmt.Errorf("error opening driver list: %s doesn't exist\nDid you run `dbc init`?", m.Path) - } else { - return fmt.Errorf("error opening driver list at %s: %w", m.Path, err) - } - } - defer f.Close() - - if err := toml.NewDecoder(f).Decode(&m.list); err != nil { return err } - if m.list.Drivers == nil { m.list.Drivers = make(map[string]driverSpec) } @@ -135,11 +121,7 @@ func (m addModel) Init() tea.Cmd { drv, err := findDriver(spec.Name, drivers) if err != nil { - // If we have registry errors, enhance the error message - if registryErrors != nil { - return fmt.Errorf("%w\n\nNote: Some driver registries were unavailable:\n%s", err, registryErrors.Error()) - } - return err + return wrapWithRegistryContext(err, registryErrors) } if spec.Vers != nil { @@ -158,9 +140,8 @@ func (m addModel) Init() tea.Cmd { // No packages. Very unlikely edge case. err = fmt.Errorf("driver `%s` not found in driver registry index", spec.Name) } - // If we have registry errors, enhance the error message if registryErrors != nil { - return fmt.Errorf("%w\n\nNote: Some driver registries were unavailable:\n%s", err, registryErrors.Error()) + return wrapWithRegistryContext(err, registryErrors) } return err } @@ -196,7 +177,7 @@ func (m addModel) Init() tea.Cmd { } } - f, err = os.Create(p) + f, err := os.Create(p) if err != nil { return fmt.Errorf("error creating file %s: %w", p, err) } diff --git a/cmd/dbc/add_test.go b/cmd/dbc/add_test.go index e72651b7..b3d2a787 100644 --- a/cmd/dbc/add_test.go +++ b/cmd/dbc/add_test.go @@ -52,7 +52,7 @@ func TestAdd(t *testing.T) { { m := AddCmd{Path: filepath.Join(dir, "dbc.toml"), Driver: []string{"test-driver-1"}}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) defer cancel() @@ -105,7 +105,7 @@ func TestAddRepeatedNewWithConstraint(t *testing.T) { { m := AddCmd{Path: filepath.Join(dir, "dbc.toml"), Driver: []string{"test-driver-1"}}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) defer cancel() @@ -129,7 +129,7 @@ func TestAddRepeatedNewWithConstraint(t *testing.T) { { m := AddCmd{Path: filepath.Join(dir, "dbc.toml"), Driver: []string{"test-driver-1>=1.0.0"}}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) defer cancel() @@ -186,7 +186,7 @@ func TestAddMultiple(t *testing.T) { { m := AddCmd{Path: filepath.Join(dir, "dbc.toml"), Driver: []string{"test-driver-2", "test-driver-1>=1.0.0"}}. GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) defer cancel() @@ -223,7 +223,7 @@ func (suite *SubcommandTestSuite) TestAddWithPre() { Driver: []string{"test-driver-2"}, Pre: true, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) @@ -248,7 +248,7 @@ func (suite *SubcommandTestSuite) TestAddWithPreOnlyPrereleaseDriver() { Driver: []string{"test-driver-only-pre"}, Pre: true, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) @@ -273,7 +273,7 @@ func (suite *SubcommandTestSuite) TestAddWithoutPreOnlyPrereleaseDriver() { Driver: []string{"test-driver-only-pre"}, Pre: false, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) out := suite.runCmdErr(m) suite.Contains(out, "driver `test-driver-only-pre` not found in driver registry index (but prerelease versions filtered out); try: dbc add --pre test-driver-only-pre") @@ -290,7 +290,7 @@ func (suite *SubcommandTestSuite) TestAddWithPreAndConstraint() { Driver: []string{"test-driver-2>=2.0.0"}, Pre: true, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) @@ -316,7 +316,7 @@ func (suite *SubcommandTestSuite) TestAddExplicitPrereleaseWithoutPreFlag() { Driver: []string{"test-driver-only-pre=0.9.0-alpha.1"}, Pre: false, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) @@ -420,7 +420,7 @@ func (suite *SubcommandTestSuite) TestAddOutput() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: []string{"test-driver-1"}, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) out := suite.runCmd(m) suite.Contains(out, "added test-driver-1 to driver list") @@ -435,7 +435,7 @@ func (suite *SubcommandTestSuite) TestAddMultipleOutput() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: []string{"test-driver-1", "test-driver-2"}, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) out := suite.runCmd(m) suite.Contains(out, "added test-driver-1 to driver list") @@ -452,7 +452,7 @@ func (suite *SubcommandTestSuite) TestAddReplacingDriverOutput() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: []string{"test-driver-1"}, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) // Add same driver with constraint and verify replacement message @@ -460,7 +460,7 @@ func (suite *SubcommandTestSuite) TestAddReplacingDriverOutput() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: []string{"test-driver-1>=1.0.0"}, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) out := suite.runCmd(m) suite.Contains(out, "replacing existing driver test-driver-1") diff --git a/cmd/dbc/auth.go b/cmd/dbc/auth.go index 7fc9e7e8..0e352b8b 100644 --- a/cmd/dbc/auth.go +++ b/cmd/dbc/auth.go @@ -28,10 +28,16 @@ import ( tea "charm.land/bubbletea/v2" "github.com/cli/browser" "github.com/cli/oauth/device" - "github.com/columnar-tech/dbc" "github.com/columnar-tech/dbc/auth" ) +func ensureHTTPS(uri string) string { + if !strings.HasPrefix(uri, "https://") { + return "https://" + uri + } + return uri +} + type AuthCmd struct { Login *LoginCmd `arg:"subcommand" help:"Authenticate with a driver registry"` Logout *LogoutCmd `arg:"subcommand" help:"Log out from a driver registry"` @@ -58,7 +64,8 @@ func (l LoginCmd) GetModelCustom(baseModel baseModel) tea.Model { reader := bufio.NewReader(os.Stdin) apiKey, err := reader.ReadString('\n') if err != nil && err != io.EOF { - panic(fmt.Errorf("failed to read API key from stdin: %w", err)) + fmt.Fprintf(os.Stderr, "failed to read API key from stdin: %s\n", err) + os.Exit(1) } l.ApiKey = strings.TrimSpace(apiKey) @@ -86,12 +93,7 @@ func (l LoginCmd) GetModelCustom(baseModel baseModel) tea.Model { } func (l LoginCmd) GetModel() tea.Model { - return l.GetModelCustom( - baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, - ) + return l.GetModelCustom(defaultBaseModel()) } type authSuccessMsg struct { @@ -113,9 +115,7 @@ type loginModel struct { } func (m loginModel) Init() tea.Cmd { - if !strings.HasPrefix(m.inputURI, "https://") { - m.inputURI = "https://" + m.inputURI - } + m.inputURI = ensureHTTPS(m.inputURI) u, err := url.Parse(m.inputURI) if err != nil { @@ -139,7 +139,7 @@ func (m loginModel) authConfig() tea.Cmd { func (m loginModel) requestDeviceCode(cfg auth.OpenIDConfig) tea.Cmd { return func() tea.Msg { - rsp, err := device.RequestCode(dbc.DefaultClient, cfg.DeviceAuthorizationEndpoint.String(), + rsp, err := device.RequestCode(dbcClient.HTTPClient(), cfg.DeviceAuthorizationEndpoint.String(), m.oauthClientID, []string{"openid", "offline_access"}) if err != nil { return fmt.Errorf("failed to request device code: %w", err) @@ -159,8 +159,8 @@ func (m loginModel) apiKeyToToken() tea.Cmd { ApiKey: m.apiKey, } - if !cred.Refresh() { - return fmt.Errorf("failed to obtain access token using provided API key") + if err := cred.Refresh(); err != nil { + return fmt.Errorf("failed to obtain access token using provided API key: %w", err) } return cred @@ -188,7 +188,7 @@ func (m loginModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { tea.Println("Opening ", msg.VerificationURIComplete, " in your default web browser..."), func() tea.Msg { browser.OpenURL(msg.VerificationURIComplete) - accessToken, err := device.Wait(context.TODO(), dbc.DefaultClient, m.tokenURI.String(), device.WaitOptions{ + accessToken, err := device.Wait(context.TODO(), dbcClient.HTTPClient(), m.tokenURI.String(), device.WaitOptions{ ClientID: m.oauthClientID, DeviceCode: msg, }) @@ -262,12 +262,7 @@ func (l LogoutCmd) GetModelCustom(baseModel baseModel) tea.Model { } func (l LogoutCmd) GetModel() tea.Model { - return l.GetModelCustom( - baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, - ) + return l.GetModelCustom(defaultBaseModel()) } type logoutModel struct { @@ -278,9 +273,7 @@ type logoutModel struct { } func (m logoutModel) Init() tea.Cmd { - if !strings.HasPrefix(m.inputURI, "https://") { - m.inputURI = "https://" + m.inputURI - } + m.inputURI = ensureHTTPS(m.inputURI) u, err := url.Parse(m.inputURI) if err != nil { @@ -326,12 +319,7 @@ func (l LicenseInstallCmd) GetModelCustom(baseModel baseModel) tea.Model { } func (l LicenseInstallCmd) GetModel() tea.Model { - return l.GetModelCustom( - baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, - ) + return l.GetModelCustom(defaultBaseModel()) } type licenseInstalledMsg struct{} @@ -369,7 +357,11 @@ func (m licenseInstallModel) FinalOutput() string { if !m.installed { return "" } - return "License installed to " + auth.LicensePath() + p, err := auth.LicensePath() + if err != nil { + return "License installed (could not determine path)" + } + return "License installed to " + p } func (m licenseInstallModel) View() tea.View { return tea.NewView("") } diff --git a/cmd/dbc/auth_test.go b/cmd/dbc/auth_test.go index 76b2ce52..246d9514 100644 --- a/cmd/dbc/auth_test.go +++ b/cmd/dbc/auth_test.go @@ -383,7 +383,9 @@ func (suite *SubcommandTestSuite) TestLicenseInstallWrongFilenameWithForce() { out := suite.runCmd(m) suite.Contains(out, "License installed") - installed, err := os.ReadFile(auth.LicensePath()) + lp, err := auth.LicensePath() + suite.Require().NoError(err) + installed, err := os.ReadFile(lp) suite.Require().NoError(err) suite.Equal("license-data", string(installed)) } @@ -433,7 +435,9 @@ func (suite *SubcommandTestSuite) TestLicenseInstallAlreadyExistsWithForce() { out := suite.runCmd(m) suite.Contains(out, "License installed") - installed, err := os.ReadFile(auth.LicensePath()) + lp, err := auth.LicensePath() + suite.Require().NoError(err) + installed, err := os.ReadFile(lp) suite.Require().NoError(err) suite.Equal("new", string(installed)) } @@ -456,7 +460,9 @@ func (suite *SubcommandTestSuite) TestLicenseInstallHappyPath() { out := suite.runCmd(m) suite.Contains(out, "License installed") - installed, err := os.ReadFile(auth.LicensePath()) + lp, err := auth.LicensePath() + suite.Require().NoError(err) + installed, err := os.ReadFile(lp) suite.Require().NoError(err) suite.Equal("license-data", string(installed)) } diff --git a/delegates.go b/cmd/dbc/delegates.go similarity index 99% rename from delegates.go rename to cmd/dbc/delegates.go index e3f6e9a2..64993eee 100644 --- a/delegates.go +++ b/cmd/dbc/delegates.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package dbc +package main import ( "fmt" diff --git a/cmd/dbc/docs.go b/cmd/dbc/docs.go index 24a0e125..5d7ab665 100644 --- a/cmd/dbc/docs.go +++ b/cmd/dbc/docs.go @@ -58,22 +58,18 @@ func (c DocsCmd) GetModelCustom(baseModel baseModel, noOpen bool, openBrowserFun } func (c DocsCmd) GetModel() tea.Model { - return c.GetModelCustom(baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, c.NoOpen, openBrowserFunc, fallbackDriverDocsUrl) + return c.GetModelCustom(defaultBaseModel(), c.NoOpen, openBrowserFunc, fallbackDriverDocsUrl) } type docsModel struct { baseModel - driver string - drv *dbc.Driver - urlToOpen string - noOpen bool - fallbackUrls map[string]string - openBrowser func(string) error - registryErrors error // Store registry errors for better error messages + driver string + drv *dbc.Driver + urlToOpen string + noOpen bool + fallbackUrls map[string]string + openBrowser func(string) error } func (m docsModel) Init() tea.Cmd { @@ -90,11 +86,7 @@ func (m docsModel) Init() tea.Cmd { drv, err := findDriver(m.driver, drivers) if err != nil { - // If we have registry errors, enhance the error message - if registryErr != nil { - return fmt.Errorf("%w\n\nNote: Some driver registries were unavailable:\n%s", err, registryErr.Error()) - } - return err + return wrapWithRegistryContext(err, registryErr) } return drv @@ -111,8 +103,8 @@ func (m docsModel) openBrowserCmd(url string) tea.Cmd { } func (m docsModel) getDocsUrlFor(driver *dbc.Driver) string { - if driver.DocsUrl != "" { - return driver.DocsUrl + if driver.DocsURL != "" { + return driver.DocsURL } fallbackUrl, keyExists := m.fallbackUrls[driver.Path] if keyExists && fallbackUrl != "" { diff --git a/cmd/dbc/driver_list.go b/cmd/dbc/driver_list.go index 8b5e6e1d..9dfc927d 100644 --- a/cmd/dbc/driver_list.go +++ b/cmd/dbc/driver_list.go @@ -44,7 +44,7 @@ func GetDriverList(fname string) ([]dbc.PkgInfo, error) { return nil, fmt.Errorf("error decoding driver list %s: %w", fname, err) } - drivers, err := dbc.GetDriverList() + drivers, err := getDriverRegistry() if err != nil { return nil, err } diff --git a/file_progress_model.go b/cmd/dbc/file_progress_model.go similarity index 96% rename from file_progress_model.go rename to cmd/dbc/file_progress_model.go index 3cb4bf45..5067696b 100644 --- a/file_progress_model.go +++ b/cmd/dbc/file_progress_model.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package dbc +package main import ( "fmt" @@ -39,6 +39,9 @@ func (m FileProgressModel) Init() tea.Cmd { func (m *FileProgressModel) SetPercent(written, total int64) tea.Cmd { m.written = written m.totalBytes = total + if total <= 0 { + return m.Model.SetPercent(0) + } return m.Model.SetPercent(float64(written) / float64(total)) } diff --git a/cmd/dbc/helpers.go b/cmd/dbc/helpers.go new file mode 100644 index 00000000..594a0089 --- /dev/null +++ b/cmd/dbc/helpers.go @@ -0,0 +1,59 @@ +// Copyright 2026 Columnar Technologies 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 main + +import ( + "errors" + "fmt" + "os" + + "github.com/pelletier/go-toml/v2" +) + +func wrapWithRegistryContext(err, registryErr error) error { + if registryErr != nil { + return fmt.Errorf("%w\n\nNote: Some driver registries were unavailable:\n%s", err, registryErr.Error()) + } + return err +} + +func defaultBaseModel() baseModel { + return baseModel{ + getDriverRegistry: getDriverRegistry, + downloadPkg: downloadPkg, + } +} + +func openAndDecodeDriverList(path string) (DriversList, error) { + p, err := driverListPath(path) + if err != nil { + return DriversList{}, err + } + + f, err := os.Open(p) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return DriversList{}, fmt.Errorf("error opening driver list: %s doesn't exist\nDid you run `dbc init`?", path) + } + return DriversList{}, fmt.Errorf("error opening driver list at %s: %w", path, err) + } + defer f.Close() + + var list DriversList + if err := toml.NewDecoder(f).Decode(&list); err != nil { + return DriversList{}, err + } + return list, nil +} diff --git a/cmd/dbc/info.go b/cmd/dbc/info.go index f1594817..3a0aeb0d 100644 --- a/cmd/dbc/info.go +++ b/cmd/dbc/info.go @@ -37,19 +37,15 @@ func (c InfoCmd) GetModelCustom(baseModel baseModel) tea.Model { } func (c InfoCmd) GetModel() tea.Model { - return c.GetModelCustom(baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }) + return c.GetModelCustom(defaultBaseModel()) } type infoModel struct { baseModel - driver string - jsonOutput bool - drv dbc.Driver - registryErrors error // Store registry errors for better error messages + driver string + jsonOutput bool + drv dbc.Driver } func (m infoModel) Init() tea.Cmd { @@ -62,11 +58,7 @@ func (m infoModel) Init() tea.Cmd { drv, err := findDriver(m.driver, drivers) if err != nil { - // If we have registry errors, enhance the error message - if registryErr != nil { - return fmt.Errorf("%w\n\nNote: Some driver registries were unavailable:\n%s", err, registryErr.Error()) - } - return err + return wrapWithRegistryContext(err, registryErr) } return drv @@ -78,7 +70,10 @@ func formatDriverInfo(drv dbc.Driver) string { return "" } - info := drv.MaxVersion() + info, ok := drv.MaxVersion() + if !ok { + return "" + } var b strings.Builder b.WriteString(bold.Render("Driver: ") + nameStyle.Render(drv.Path) + "\n") @@ -88,14 +83,17 @@ func formatDriverInfo(drv dbc.Driver) string { b.WriteString(bold.Render("Description: ") + drv.Desc + "\n") b.WriteString(bold.Render("Available Packages:") + "\n") for _, pkg := range info.Packages { - b.WriteString(" - " + descStyle.Render(pkg.PlatformTuple) + "\n") + b.WriteString(" - " + descStyle.Render(pkg.Platform) + "\n") } return strings.TrimSuffix(b.String(), "\n") } func driverInfoJSON(drv dbc.Driver) string { - info := drv.MaxVersion() + info, ok := drv.MaxVersion() + if !ok { + return "{}" + } var driverInfoOutput = struct { Driver string `json:"driver"` @@ -112,7 +110,7 @@ func driverInfoJSON(drv dbc.Driver) string { Desc: drv.Desc, } for _, pkg := range info.Packages { - driverInfoOutput.Packages = append(driverInfoOutput.Packages, pkg.PlatformTuple) + driverInfoOutput.Packages = append(driverInfoOutput.Packages, pkg.Platform) } jsonBytes, err := json.Marshal(driverInfoOutput) diff --git a/cmd/dbc/info_test.go b/cmd/dbc/info_test.go index c3fc28f4..a6388985 100644 --- a/cmd/dbc/info_test.go +++ b/cmd/dbc/info_test.go @@ -22,7 +22,7 @@ import ( func (suite *SubcommandTestSuite) TestInfo() { m := InfoCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("\r ", "Driver: test-driver-1\n"+ @@ -35,7 +35,7 @@ func (suite *SubcommandTestSuite) TestInfo() { func (suite *SubcommandTestSuite) TestInfo_DriverNotFound() { m := InfoCmd{Driver: "non-existent-driver"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmdErr(m) suite.validateOutput("\r ", "\nError: driver `non-existent-driver` not found in driver registry index; try: `dbc search` to list available drivers", out) diff --git a/cmd/dbc/install.go b/cmd/dbc/install.go index 592a21a2..d4d61221 100644 --- a/cmd/dbc/install.go +++ b/cmd/dbc/install.go @@ -32,6 +32,17 @@ import ( "github.com/columnar-tech/dbc/config" ) +func manifestToPackageInfo(m config.Manifest) dbc.PkgInfo { + return dbc.PkgInfo{ + Driver: dbc.Driver{ + Title: m.Name, + Path: m.ID, + License: m.License, + }, + Version: m.Version, + } +} + func parseDriverConstraint(driver string) (string, *semver.Constraints, error) { driver = strings.TrimSpace(driver) splitIdx := strings.IndexAny(driver, " ~^<>=!") @@ -81,7 +92,7 @@ func (c InstallCmd) GetModelCustom(baseModel baseModel) tea.Model { baseModel: baseModel, isLocal: isLocal, localPackagePath: localPackagePath, - p: dbc.NewFileProgress( + p: NewFileProgress( progress.WithDefaultBlend(), progress.WithWidth(20), progress.WithoutPercentage(), @@ -90,10 +101,7 @@ func (c InstallCmd) GetModelCustom(baseModel baseModel) tea.Model { } func (c InstallCmd) GetModel() tea.Model { - return c.GetModelCustom(baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }) + return c.GetModelCustom(defaultBaseModel()) } func verifySignature(m config.Manifest, noVerify bool) error { @@ -179,7 +187,7 @@ type progressiveInstallModel struct { state installState spinner spinner.Model - p dbc.FileProgressModel + p FileProgressModel width, height int isLocal bool @@ -218,16 +226,23 @@ func (m progressiveInstallModel) Preamble() string { return "" } +func (m progressiveInstallModel) hasConflict() bool { + return m.conflictingInfo.ID != "" && m.conflictingInfo.Version != nil +} + +func (m progressiveInstallModel) isAlreadyInstalled() bool { + return m.conflictingInfo.ID != "" && m.conflictingInfo.Version != nil && + m.conflictingInfo.Version.Equal(m.DriverPackage.Version) +} + func (m progressiveInstallModel) FinalOutput() string { - if m.conflictingInfo.ID != "" && m.conflictingInfo.Version != nil { - if m.conflictingInfo.Version.Equal(m.DriverPackage.Version) { - if m.jsonOutput { - return fmt.Sprintf(`{"status":"already installed","driver":"%s","version":"%s","location":"%s"}`, - m.conflictingInfo.ID, m.conflictingInfo.Version, filepath.SplitList(m.cfg.Location)[0]) - } - return fmt.Sprintf("\nDriver %s %s already installed at %s", + if m.isAlreadyInstalled() { + if m.jsonOutput { + return fmt.Sprintf(`{"status":"already installed","driver":"%s","version":"%s","location":"%s"}`, m.conflictingInfo.ID, m.conflictingInfo.Version, filepath.SplitList(m.cfg.Location)[0]) } + return fmt.Sprintf("\nDriver %s %s already installed at %s", + m.conflictingInfo.ID, m.conflictingInfo.Version, filepath.SplitList(m.cfg.Location)[0]) } var b strings.Builder @@ -245,7 +260,7 @@ func (m progressiveInstallModel) FinalOutput() string { output.Driver = m.Driver output.Version = m.DriverPackage.Version.String() output.Location = filepath.SplitList(m.cfg.Location)[0] - if m.conflictingInfo.ID != "" && m.conflictingInfo.Version != nil { + if m.hasConflict() { output.Conflict = fmt.Sprintf("%s (version: %s)", m.conflictingInfo.ID, m.conflictingInfo.Version) } @@ -319,11 +334,9 @@ func (m progressiveInstallModel) searchForDriver(list []dbc.Driver) (tea.Model, func (m progressiveInstallModel) startDownloading() (tea.Model, tea.Cmd) { m.state = stDownloading - if m.conflictingInfo.ID != "" && m.conflictingInfo.Version != nil { - if m.conflictingInfo.Version.Equal(m.DriverPackage.Version) { - m.state = stDone - return m, tea.Quit - } + if m.isAlreadyInstalled() { + m.state = stDone + return m, tea.Quit } return m, func() tea.Msg { @@ -408,7 +421,7 @@ func (m progressiveInstallModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m.startInstalling(msg) case config.Manifest: if m.DriverPackage.Version == nil { - m.DriverPackage = msg.ToPackageInfo() + m.DriverPackage = manifestToPackageInfo(msg) } m.state = stVerifying @@ -451,10 +464,8 @@ func (m progressiveInstallModel) View() tea.View { return tea.NewView("") } - if m.conflictingInfo.ID != "" && m.conflictingInfo.Version != nil { - if m.conflictingInfo.Version.Equal(m.DriverPackage.Version) { - return tea.NewView("") - } + if m.isAlreadyInstalled() { + return tea.NewView("") } var b strings.Builder diff --git a/cmd/dbc/install_test.go b/cmd/dbc/install_test.go index c8048771..1db1ee9a 100644 --- a/cmd/dbc/install_test.go +++ b/cmd/dbc/install_test.go @@ -28,7 +28,7 @@ import ( func (suite *SubcommandTestSuite) TestInstall() { m := InstallCmd{Driver: "test-driver-1", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", @@ -38,7 +38,7 @@ func (suite *SubcommandTestSuite) TestInstall() { func (suite *SubcommandTestSuite) TestInstallDriverNotFound() { m := InstallCmd{Driver: "foo", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r ", "\nError: could not find driver: driver `foo` not found in driver registry index; try: `dbc search` to list available drivers", suite.runCmdErr(m)) suite.driverIsNotInstalled("test-driver-1") } @@ -58,14 +58,14 @@ func (suite *SubcommandTestSuite) TestInstallWithVersion() { for _, tt := range tests { suite.Run(tt.driver, func() { m := InstallCmd{Driver: tt.driver, Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nInstalled test-driver-1 "+tt.expectedVersion+" to "+suite.Dir(), out) suite.driverIsInstalled("test-driver-1", true) m = UninstallCmd{Driver: "test-driver-1", Level: suite.configLevel}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) }) } @@ -73,7 +73,7 @@ func (suite *SubcommandTestSuite) TestInstallWithVersion() { func (suite *SubcommandTestSuite) TestInstallWithVersionLessSpace() { m := InstallCmd{Driver: "test-driver-1 < 1.1.0"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nInstalled test-driver-1 1.0.0 to "+suite.tempdir, out) @@ -81,12 +81,12 @@ func (suite *SubcommandTestSuite) TestInstallWithVersionLessSpace() { func (suite *SubcommandTestSuite) TestReinstallUpdateVersion() { m := InstallCmd{Driver: "test-driver-1<=1.0.0"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nInstalled test-driver-1 1.0.0 to "+suite.tempdir, suite.runCmd(m)) m = InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nRemoved conflicting driver: test-driver-1 (version: 1.0.0)\nInstalled test-driver-1 1.1.0 to "+suite.tempdir, suite.runCmd(m)) @@ -100,7 +100,7 @@ func (suite *SubcommandTestSuite) TestInstallVenv() { suite.T().Setenv("VIRTUAL_ENV", suite.tempdir) m := InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nInstalled test-driver-1 1.1.0 to "+filepath.Join(suite.tempdir, "etc", "adbc", "drivers"), suite.runCmd(m)) } @@ -118,7 +118,7 @@ func (suite *SubcommandTestSuite) TestInstallEnvironmentPrecedence() { suite.T().Setenv("CONDA_PREFIX", conda_path) m := InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.FileExists(filepath.Join(driver_path, "test-driver-1.toml")) @@ -127,14 +127,14 @@ func (suite *SubcommandTestSuite) TestInstallEnvironmentPrecedence() { suite.T().Setenv("ADBC_DRIVER_PATH", "") m = InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.FileExists(filepath.Join(venv_path, "etc", "adbc", "drivers", "test-driver-1.toml")) suite.NoFileExists(filepath.Join(conda_path, "etc", "adbc", "drivers", "test-driver-1.toml")) suite.T().Setenv("VIRTUAL_ENV", "") m = InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.FileExists(filepath.Join(conda_path, "etc", "adbc", "drivers", "test-driver-1.toml")) } @@ -144,14 +144,14 @@ func (suite *SubcommandTestSuite) TestInstallCondaPrefix() { suite.T().Setenv("CONDA_PREFIX", suite.tempdir) m := InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nInstalled test-driver-1 1.1.0 to "+filepath.Join(suite.tempdir, "etc", "adbc", "drivers"), suite.runCmd(m)) } func (suite *SubcommandTestSuite) TestInstallManifestOnlyDriver() { m := InstallCmd{Driver: "test-driver-manifest-only", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nInstalled test-driver-manifest-only 1.0.0 to "+suite.Dir()+ @@ -161,7 +161,7 @@ func (suite *SubcommandTestSuite) TestInstallManifestOnlyDriver() { func (suite *SubcommandTestSuite) TestInstallDriverNoSignature() { m := InstallCmd{Driver: "test-driver-no-sig"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmdErr(m) suite.Contains(out, "signature file 'test-driver-1-not-valid.so.sig' for driver is missing") @@ -171,7 +171,7 @@ func (suite *SubcommandTestSuite) TestInstallDriverNoSignature() { // Note: The UI output (first parameter) serves as documentation but isn't verified // by validateOutput due to tea.WithoutRenderer() mode. Manual verification needed. m = InstallCmd{Driver: "test-driver-no-sig", NoVerify: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[-] verifying signature\r\n", "\nInstalled test-driver-no-sig 1.0.0 to "+suite.tempdir, suite.runCmd(m)) } @@ -184,7 +184,7 @@ func (suite *SubcommandTestSuite) TestInstallGitignoreDefaultBehavior() { suite.NoFileExists(ignorePath) m := InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) _ = suite.runCmd(m) suite.FileExists(ignorePath) @@ -205,7 +205,7 @@ func (suite *SubcommandTestSuite) TestInstallGitignoreExistingDir() { suite.NoFileExists(ignorePath) m := InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) _ = suite.runCmd(m) // There shouldn't be a .gitignore because we didn't create the dir fresh @@ -222,7 +222,7 @@ func (suite *SubcommandTestSuite) TestInstallGitignorePreserveUserModified() { // First install - should create .gitignore m := InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) _ = suite.runCmd(m) suite.FileExists(ignorePath) @@ -236,10 +236,10 @@ func (suite *SubcommandTestSuite) TestInstallGitignorePreserveUserModified() { // Second install - should preserve user's modifications m = UninstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) _ = suite.runCmd(m) m = InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) _ = suite.runCmd(m) // Verify the user's content is preserved @@ -257,7 +257,7 @@ func (suite *SubcommandTestSuite) TestInstallCreatesSymlinks() { // Install a driver m := InstallCmd{Driver: "test-driver-1", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) _ = suite.runCmd(m) suite.driverIsInstalled("test-driver-1", true) @@ -272,7 +272,7 @@ func (suite *SubcommandTestSuite) TestInstallCreatesSymlinks() { func (suite *SubcommandTestSuite) TestInstallLocalPackage() { packagePath := filepath.Join("testdata", "test-driver-1.tar.gz") m := InstallCmd{Driver: packagePath, Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("Installing from local package: "+packagePath+"\r\n\r\n\r"+ @@ -284,7 +284,7 @@ func (suite *SubcommandTestSuite) TestInstallLocalPackage() { func (suite *SubcommandTestSuite) TestInstallLocalPackageNotFound() { packagePath := filepath.Join("testdata", "test-driver-2.tar.gz") m := InstallCmd{Driver: packagePath, Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmdErr(m) errmsg := "no such file or directory" @@ -299,7 +299,7 @@ func (suite *SubcommandTestSuite) TestInstallLocalPackageNotFound() { func (suite *SubcommandTestSuite) TestInstallLocalPackageNoSignature() { packagePath := filepath.Join("testdata", "test-driver-no-sig.tar.gz") m := InstallCmd{Driver: packagePath}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmdErr(m) suite.Contains(out, "signature file 'test-driver-1-not-valid.so.sig' for driver is missing") @@ -307,7 +307,7 @@ func (suite *SubcommandTestSuite) TestInstallLocalPackageNoSignature() { suite.NoDirExists(filepath.Join(suite.tempdir, "test-driver-no-sig")) m = InstallCmd{Driver: packagePath, NoVerify: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("Installing from local package: "+packagePath+"\r\n\r\n\r"+ "[✓] installing\r\n[-] verifying signature\r\n", "\nInstalled test-driver-no-sig 1.1.0 to "+suite.tempdir, suite.runCmd(m)) @@ -319,7 +319,7 @@ func (suite *SubcommandTestSuite) TestInstallLocalPackageFixUpName() { packagePath := filepath.Join(suite.tempdir, "test-driver-1_"+config.PlatformTuple()+"_v1.0.0.tgz") suite.Require().NoError(os.Symlink(origPackagePath, packagePath)) m := InstallCmd{Driver: packagePath, Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("Installing from local package: "+packagePath+"\r\n\r\n\r"+ @@ -331,7 +331,7 @@ func (suite *SubcommandTestSuite) TestInstallLocalPackageFixUpName() { func (suite *SubcommandTestSuite) TestInstallWithPreOnlyPrereleaseDriver() { // Install test-driver-only-pre with --pre flag, should succeed m := InstallCmd{Driver: "test-driver-only-pre", Level: suite.configLevel, Pre: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", @@ -342,7 +342,7 @@ func (suite *SubcommandTestSuite) TestInstallWithPreOnlyPrereleaseDriver() { func (suite *SubcommandTestSuite) TestInstallWithoutPreOnlyPrereleaseDriver() { // Try to install test-driver-only-pre without --pre flag, should fail m := InstallCmd{Driver: "test-driver-only-pre", Level: suite.configLevel, Pre: false}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmdErr(m) suite.Contains(out, "driver `test-driver-only-pre` not found") @@ -353,12 +353,12 @@ func (suite *SubcommandTestSuite) TestInstallWithoutPreOnlyPrereleaseDriver() { func (suite *SubcommandTestSuite) TestInstallWithoutPreWhenPrereleaseAlreadyInstalled() { m := InstallCmd{Driver: "test-driver-only-pre", Level: suite.configLevel, Pre: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.driverIsInstalled("test-driver-only-pre", false) m = InstallCmd{Driver: "test-driver-only-pre", Level: suite.configLevel, Pre: false}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmdErr(m) suite.Contains(out, "already installed") @@ -369,7 +369,7 @@ func (suite *SubcommandTestSuite) TestInstallWithoutPreWhenPrereleaseAlreadyInst func (suite *SubcommandTestSuite) TestInstallExplicitPrereleaseWithoutPreFlag() { // Install explicit prerelease version WITHOUT --pre flag, should succeed per requirement m := InstallCmd{Driver: "test-driver-only-pre=0.9.0-alpha.1", Level: suite.configLevel, Pre: false}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmd(m) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", @@ -456,7 +456,7 @@ func (suite *SubcommandTestSuite) TestInstallDriverWithSubdirectories() { // Should fail m := InstallCmd{Driver: packagePath, NoVerify: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := suite.runCmdErr(m) // and return an error with this diff --git a/cmd/dbc/main.go b/cmd/dbc/main.go index f6f4bccb..43386a0d 100644 --- a/cmd/dbc/main.go +++ b/cmd/dbc/main.go @@ -20,6 +20,7 @@ import ( "os" "slices" "strings" + "sync" tea "charm.land/bubbletea/v2" "charm.land/lipgloss/v2" @@ -78,8 +79,34 @@ type NeedsRenderer interface { NeedsRenderer() } +var ( + dbcClient *dbc.Client + dbcClientErr error + dbcClientOnce sync.Once +) + +func newDefaultClient() (*dbc.Client, error) { + var opts []dbc.Option + if val := os.Getenv("DBC_BASE_URL"); val != "" { + opts = append(opts, dbc.WithBaseURL(val)) + } + return dbc.NewClient(opts...) +} + +func initDBCClient() error { + dbcClientOnce.Do(func() { + dbcClient, dbcClientErr = newDefaultClient() + }) + return dbcClientErr +} + // use this so we can override this in tests -var getDriverRegistry = dbc.GetDriverList +var getDriverRegistry = func() ([]dbc.Driver, error) { + if err := initDBCClient(); err != nil { + return nil, fmt.Errorf("failed to initialize client: %w", err) + } + return dbcClient.Search("") +} func findDriver(name string, drivers []dbc.Driver) (dbc.Driver, error) { idx := slices.IndexFunc(drivers, func(d dbc.Driver) bool { @@ -217,9 +244,14 @@ func main() { args cmds ) + if err := initDBCClient(); err != nil { + fmt.Fprintf(os.Stderr, "error initializing client: %v\n", err) + os.Exit(1) + } + p, err := newParser(&args) if err != nil { - fmt.Println("Error creating argument parser:", err) + fmt.Fprintf(os.Stderr, "error creating argument parser: %v\n", err) os.Exit(1) } @@ -261,6 +293,9 @@ func main() { os.Exit(0) case modelCmd: m = sub.GetModel() + default: + fmt.Fprintf(os.Stderr, "internal error: unrecognized subcommand %T\n", p.Subcommand()) + os.Exit(1) } // f, err := tea.LogToFile("debug.log", "debug") diff --git a/cmd/dbc/main_test.go b/cmd/dbc/main_test.go index 2322a8ee..e674a0d0 100644 --- a/cmd/dbc/main_test.go +++ b/cmd/dbc/main_test.go @@ -110,11 +110,14 @@ func TestCmdStatus(t *testing.T) { var m tea.Model var err error - go func() { m, err = p.Run() }() + done := make(chan struct{}) + go func() { + m, err = p.Run() + close(done) + }() - <-time.After(time.Second * 1) + <-done - p.Wait() require.NoError(t, err, out.String()) if h, ok := m.(HasStatus); ok { diff --git a/cmd/dbc/registry_test.go b/cmd/dbc/registry_test.go index 33e04661..d4233dee 100644 --- a/cmd/dbc/registry_test.go +++ b/cmd/dbc/registry_test.go @@ -96,7 +96,7 @@ func (s *RegistryTestSuite) TearDownTest() { func (s *RegistryTestSuite) TestInstallDriver() { m := InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := s.run(m) s.Equal("\nInstalled test-driver-1 1.1.0 to "+s.cfgUserPath, out) @@ -117,7 +117,7 @@ func (s *RegistryTestSuite) TestInstallDriver() { func (s *RegistryTestSuite) TestPartialReinstallDriver() { // First install the driver normally. m := InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out := s.run(m) s.Equal("\nInstalled test-driver-1 1.1.0 to "+s.cfgUserPath, out) @@ -125,7 +125,7 @@ func (s *RegistryTestSuite) TestPartialReinstallDriver() { // Now reinstall the driver, which should succeed even though the registry key is missing. m = InstallCmd{Driver: "test-driver-1"}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) out = s.run(m) s.Equal("\nInstalled test-driver-1 1.1.0 to "+s.cfgUserPath, out) } diff --git a/cmd/dbc/remove.go b/cmd/dbc/remove.go index d71e4419..6baf2ce5 100644 --- a/cmd/dbc/remove.go +++ b/cmd/dbc/remove.go @@ -15,7 +15,6 @@ package main import ( - "errors" "fmt" "os" "strings" @@ -39,12 +38,9 @@ func (c RemoveCmd) GetModelCustom(baseModel baseModel) tea.Model { func (c RemoveCmd) GetModel() tea.Model { return removeModel{ - Driver: c.Driver, - Path: c.Path, - baseModel: baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, + Driver: c.Driver, + Path: c.Path, + baseModel: defaultBaseModel(), } } @@ -65,17 +61,8 @@ func (m removeModel) Init() tea.Cmd { return err } - f, err := os.Open(p) + m.list, err = openAndDecodeDriverList(m.Path) if err != nil { - if errors.Is(err, os.ErrNotExist) { - return fmt.Errorf("error opening driver list: %s doesn't exist\nDid you run `dbc init`?", m.Path) - } else { - return fmt.Errorf("error opening driver list at %s: %w", m.Path, err) - } - } - defer f.Close() - - if err := toml.NewDecoder(f).Decode(&m.list); err != nil { return err } @@ -91,7 +78,7 @@ func (m removeModel) Init() tea.Cmd { delete(m.list.Drivers, m.Driver) - f, err = os.Create(p) + f, err := os.Create(p) if err != nil { return fmt.Errorf("error creating file %s: %w", p, err) } diff --git a/cmd/dbc/remove_test.go b/cmd/dbc/remove_test.go index f8c0513d..a31c0d4e 100644 --- a/cmd/dbc/remove_test.go +++ b/cmd/dbc/remove_test.go @@ -27,7 +27,7 @@ func (suite *SubcommandTestSuite) TestRemoveOutput() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: []string{"test-driver-1"}, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) // Remove the driver and verify output @@ -35,7 +35,7 @@ func (suite *SubcommandTestSuite) TestRemoveOutput() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: "test-driver-1", }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) out := suite.runCmd(m) suite.Contains(out, "removed 'test-driver-1' from driver list") @@ -50,7 +50,7 @@ func (suite *SubcommandTestSuite) TestRemoveNonexistentDriverError() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: []string{"test-driver-1"}, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) // Try to remove a driver that doesn't exist @@ -58,7 +58,7 @@ func (suite *SubcommandTestSuite) TestRemoveNonexistentDriverError() { Path: filepath.Join(suite.tempdir, "dbc.toml"), Driver: "nonexistent-driver", }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) out := suite.runCmdErr(m) suite.Contains(out, "driver 'nonexistent-driver' not found") @@ -70,7 +70,7 @@ func (suite *SubcommandTestSuite) TestRemoveFromNonexistentFile() { Path: filepath.Join(suite.tempdir, "nonexistent-dbc.toml"), Driver: "test-driver-1", }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) out := suite.runCmdErr(m) suite.Contains(out, "doesn't exist") diff --git a/cmd/dbc/search.go b/cmd/dbc/search.go index bf0a10a0..dd046e41 100644 --- a/cmd/dbc/search.go +++ b/cmd/dbc/search.go @@ -54,10 +54,7 @@ func (s SearchCmd) GetModelCustom(baseModel baseModel) tea.Model { } func (s SearchCmd) GetModel() tea.Model { - return s.GetModelCustom(baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }) + return s.GetModelCustom(defaultBaseModel()) } type searchModel struct { diff --git a/cmd/dbc/search_test.go b/cmd/dbc/search_test.go index f0376d2c..2f578e89 100644 --- a/cmd/dbc/search_test.go +++ b/cmd/dbc/search_test.go @@ -27,8 +27,7 @@ import ( func (suite *SubcommandTestSuite) TestSearchCmd() { m := SearchCmd{}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "test-driver-1 This is a test driver \n"+ "test-driver-2 This is another test driver \n"+ @@ -40,12 +39,11 @@ func (suite *SubcommandTestSuite) TestSearchCmd() { func (suite *SubcommandTestSuite) TestSearchCmdWithInstalled() { m := InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) m = SearchCmd{}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "test-driver-1 This is a test driver [installed: env=>1.1.0]\n"+ "test-driver-2 This is another test driver \n"+ @@ -57,8 +55,7 @@ func (suite *SubcommandTestSuite) TestSearchCmdWithInstalled() { func (suite *SubcommandTestSuite) TestSearchCmdVerbose() { m := SearchCmd{Verbose: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "• test-driver-1\n Title: Test Driver 1\n "+ "Description: This is a test driver\n License: MIT\n "+ "Available Versions:\n ├── 1.0.0\n ╰── 1.1.0\n"+ @@ -81,12 +78,11 @@ func (suite *SubcommandTestSuite) TestSearchCmdVerbose() { func (suite *SubcommandTestSuite) TestSearchCmdVerboseWithInstalled() { m := InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) m = SearchCmd{Verbose: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "• test-driver-1\n Title: Test Driver 1\n "+ "Description: This is a test driver\n License: MIT\n "+ "Installed Versions:\n ╰── 1.1.0\n ╰── env => "+filepath.Join(suite.tempdir)+ @@ -123,7 +119,7 @@ func (suite *SubcommandTestSuite) TestSearchCmdVerboseWithInstalled() { func (suite *SubcommandTestSuite) TestSearchCmdWithMissingVersionInManifest() { // Install a driver m := InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) // Corrupt the manifest by removing the version key @@ -145,16 +141,14 @@ func (suite *SubcommandTestSuite) TestSearchCmdWithMissingVersionInManifest() { suite.Require().NotPanics(func() { m = SearchCmd{}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.runCmd(m) }, "Search should not panic when manifest is missing version key") } func (suite *SubcommandTestSuite) TestSearchCmdWithPre() { m := SearchCmd{Pre: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "test-driver-1 This is a test driver \n"+ "test-driver-2 This is another test driver \n"+ @@ -167,8 +161,7 @@ func (suite *SubcommandTestSuite) TestSearchCmdWithPre() { func (suite *SubcommandTestSuite) TestSearchCmdVerboseWithPre() { m := SearchCmd{Verbose: true, Pre: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "• test-driver-1\n Title: Test Driver 1\n "+ "Description: This is a test driver\n License: MIT\n "+ "Available Versions:\n ├── 1.0.0\n ╰── 1.1.0\n"+ @@ -194,12 +187,11 @@ func (suite *SubcommandTestSuite) TestSearchCmdVerboseWithPre() { func (suite *SubcommandTestSuite) TestSearchCmdWithInstalledPre() { m := InstallCmd{Driver: "test-driver-only-pre", Level: config.ConfigEnv, Pre: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) m = SearchCmd{}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "test-driver-1 This is a test driver \n"+ "test-driver-2 This is another test driver \n"+ @@ -212,18 +204,16 @@ func (suite *SubcommandTestSuite) TestSearchCmdWithInstalledPre() { func (suite *SubcommandTestSuite) TestSearchCmdVerboseWithInstalledPre() { searchModel := SearchCmd{Verbose: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) before := suite.runCmd(searchModel) suite.NotContains(before, "test-driver-only-pre") m := InstallCmd{Driver: "test-driver-only-pre", Level: config.ConfigEnv, Pre: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) searchModel = SearchCmd{Verbose: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) after := suite.runCmd(searchModel) suite.Contains(after, "test-driver-only-pre") @@ -236,18 +226,16 @@ func (suite *SubcommandTestSuite) TestSearchCmdVerboseWithInstalledPre() { func (suite *SubcommandTestSuite) TestSearchCmdNonVerboseWithInstalledPre() { searchModel := SearchCmd{}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) before := suite.runCmd(searchModel) suite.NotContains(before, "test-driver-only-pre") m := InstallCmd{Driver: "test-driver-only-pre", Level: config.ConfigEnv, Pre: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) searchModel = SearchCmd{}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) after := suite.runCmd(searchModel) suite.Contains(after, "test-driver-only-pre") @@ -258,18 +246,16 @@ func (suite *SubcommandTestSuite) TestSearchCmdNonVerboseWithInstalledPre() { func (suite *SubcommandTestSuite) TestSearchCmdWithInstalledPreJSON() { searchModel := SearchCmd{Json: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) before := suite.runCmd(searchModel) suite.NotContains(before, `"test-driver-only-pre"`) m := InstallCmd{Driver: "test-driver-only-pre", Level: config.ConfigEnv, Pre: true}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) searchModel = SearchCmd{Json: true}.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, - downloadPkg: downloadTestPkg}) + testBaseModel()) after := suite.runCmd(searchModel) suite.Contains(after, `"test-driver-only-pre"`) diff --git a/cmd/dbc/subcommand_test.go b/cmd/dbc/subcommand_test.go index 128e2e45..4b8cd507 100644 --- a/cmd/dbc/subcommand_test.go +++ b/cmd/dbc/subcommand_test.go @@ -58,6 +58,10 @@ func getTestDriverRegistry() ([]dbc.Driver, error) { return drivers.Drivers, nil } +func testBaseModel() baseModel { + return baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg} +} + func downloadTestPkg(pkg dbc.PkgInfo) (*os.File, error) { switch pkg.Driver.Path { case "test-driver-1": diff --git a/cmd/dbc/sync.go b/cmd/dbc/sync.go index ea61a684..1af29b94 100644 --- a/cmd/dbc/sync.go +++ b/cmd/dbc/sync.go @@ -50,13 +50,10 @@ func (c SyncCmd) GetModelCustom(baseModel baseModel) tea.Model { func (c SyncCmd) GetModel() tea.Model { return syncModel{ - Path: c.Path, - cfg: getConfig(c.Level), - NoVerify: c.NoVerify, - baseModel: baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, + Path: c.Path, + cfg: getConfig(c.Level), + NoVerify: c.NoVerify, + baseModel: defaultBaseModel(), } } @@ -118,23 +115,10 @@ func (s syncModel) Init() tea.Cmd { } func loadDriverList(path string) (DriversList, error) { - f, err := os.Open(path) + list, err := openAndDecodeDriverList(path) if err != nil { - var outError error - if errors.Is(err, os.ErrNotExist) { - outError = fmt.Errorf("error opening driver list: %s doesn't exist\ndid you run `dbc init`?", path) - } else { - outError = fmt.Errorf("error opening driver list at %s: %w", path, err) - } - return DriversList{}, outError - } - defer f.Close() - - var list DriversList - if err := toml.NewDecoder(f).Decode(&list); err != nil { return DriversList{}, err } - if len(list.Drivers) == 0 { return DriversList{}, fmt.Errorf("no drivers found in driver list `%s`", path) } @@ -165,11 +149,7 @@ func (s syncModel) createInstallList(list DriversList) ([]installItem, error) { // locate the driver info in the CDN driver registry index drv, err := findDriver(name, s.driverIndex) if err != nil { - // If we have registry errors, enhance the error message - if s.registryErrors != nil { - return nil, fmt.Errorf("%w\n\nNote: Some driver registries were unavailable:\n%s", err, s.registryErrors.Error()) - } - return nil, err + return nil, wrapWithRegistryContext(err, s.registryErrors) } var pkg dbc.PkgInfo diff --git a/cmd/dbc/sync_test.go b/cmd/dbc/sync_test.go index bb33e839..b91c7de0 100644 --- a/cmd/dbc/sync_test.go +++ b/cmd/dbc/sync_test.go @@ -32,14 +32,14 @@ func (suite *SubcommandTestSuite) TestSync() { m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-1-1.1.0\r\n\rDone!\r\n", "", suite.runCmd(m)) suite.FileExists(filepath.Join(suite.tempdir, "test-driver-1.toml")) m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-1-1.1.0 already installed\r\n\rDone!\r\n", "", suite.runCmd(m)) } @@ -66,7 +66,7 @@ func (suite *SubcommandTestSuite) TestSyncWithVersion() { m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-1-"+tt.expectedVersion+"\r\n\rDone!\r\n", "", suite.runCmd(m)) suite.FileExists(filepath.Join(suite.tempdir, "test-driver-1.toml")) suite.FileExists(filepath.Join(suite.tempdir, "dbc.lock")) @@ -92,14 +92,14 @@ func (suite *SubcommandTestSuite) TestSyncVirtualEnv() { m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-1-1.1.0\r\n\rDone!\r\n", "", suite.runCmd(m)) suite.FileExists(filepath.Join(suite.tempdir, "etc", "adbc", "drivers", "test-driver-1.toml")) m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-1-1.1.0 already installed\r\n\rDone!\r\n", "", suite.runCmd(m)) } @@ -117,14 +117,14 @@ func (suite *SubcommandTestSuite) TestSyncCondaPrefix() { m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-1-1.1.0\r\n\rDone!\r\n", "", suite.runCmd(m)) suite.FileExists(filepath.Join(suite.tempdir, "etc", "adbc", "drivers", "test-driver-1.toml")) m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-1-1.1.0 already installed\r\n\rDone!\r\n", "", suite.runCmd(m)) } @@ -138,7 +138,7 @@ func (suite *SubcommandTestSuite) TestSyncInstallFailSig() { m = SyncCmd{ Path: filepath.Join(suite.tempdir, "dbc.toml"), }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("\r ", "\nError: failed to verify signature: signature file 'test-driver-1-not-valid.so.sig' for driver is missing", suite.runCmdErr(m)) @@ -156,7 +156,7 @@ func (suite *SubcommandTestSuite) TestSyncInstallNoVerify() { Path: filepath.Join(suite.tempdir, "dbc.toml"), NoVerify: true, }.GetModelCustom( - baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + testBaseModel()) suite.validateOutput("✓ test-driver-no-sig-1.1.0\r\n\rDone!\r\n", "", suite.runCmd(m)) } diff --git a/cmd/dbc/tui_driver_list.go b/cmd/dbc/tui_driver_list.go index 4ffe93de..fff777d1 100644 --- a/cmd/dbc/tui_driver_list.go +++ b/cmd/dbc/tui_driver_list.go @@ -18,21 +18,17 @@ import ( "fmt" "os" - "github.com/Masterminds/semver/v3" "charm.land/bubbles/v2/list" tea "charm.land/bubbletea/v2" "charm.land/lipgloss/v2" + "github.com/Masterminds/semver/v3" "github.com/columnar-tech/dbc" "github.com/columnar-tech/dbc/config" ) const defaultWidth = 40 -var ( - docStyle = lipgloss.NewStyle().Margin(1, 2) - itemStyle = lipgloss.NewStyle().PaddingLeft(4) - selectedItemStyle = lipgloss.NewStyle().PaddingLeft(2).Foreground(lipgloss.Color("170")) -) +var docStyle = lipgloss.NewStyle().Margin(1, 2) type item struct { d dbc.Driver @@ -53,7 +49,7 @@ type model struct { } func getDrivers() tea.Msg { - drivers, err := dbc.GetDriverList() + drivers, err := getDriverRegistry() if err != nil { fmt.Println("Error getting drivers:", err) os.Exit(1) @@ -91,7 +87,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } m.chooseVersion = versionModel{ - list: list.New(versions, dbc.SimpleItemDelegate{Prompt: ">"}, 40, 15), + list: list.New(versions, SimpleItemDelegate{Prompt: ">"}, 40, 15), choice: "", } m.chooseVersion.list.Title = fmt.Sprintf("Versions for %s", i.d.Title) @@ -125,7 +121,7 @@ func (m model) View() tea.View { type versionOption semver.Version func (v versionOption) FilterValue() string { return v.String() } -func (v versionOption) String() string { return v.String() } +func (v versionOption) String() string { return semver.Version(v).String() } type versionModel struct { list list.Model diff --git a/cmd/dbc/tui_menu.go b/cmd/dbc/tui_menu.go index 74abf210..8bdd67b8 100644 --- a/cmd/dbc/tui_menu.go +++ b/cmd/dbc/tui_menu.go @@ -21,7 +21,6 @@ import ( "charm.land/bubbles/v2/list" tea "charm.land/bubbletea/v2" - "github.com/columnar-tech/dbc/config" ) func getTuiModel() tea.Model { @@ -39,7 +38,7 @@ func getTuiModel() tea.Model { driversModel.list.SetFilteringEnabled(true) m.options.SetItems([]list.Item{ - menuOption{title: "Current Config", delegate: config.Model{Prev: &m}}, + menuOption{title: "Current Config", delegate: configViewModel{Prev: &m}}, menuOption{title: "Drivers", delegate: &driversModel}, }) diff --git a/cmd/dbc/uninstall.go b/cmd/dbc/uninstall.go index 9de84122..feb61905 100644 --- a/cmd/dbc/uninstall.go +++ b/cmd/dbc/uninstall.go @@ -40,10 +40,7 @@ func (c UninstallCmd) GetModelCustom(baseModel baseModel) tea.Model { func (c UninstallCmd) GetModel() tea.Model { return uninstallModel{ - baseModel: baseModel{ - getDriverRegistry: getDriverRegistry, - downloadPkg: downloadPkg, - }, + baseModel: defaultBaseModel(), Driver: c.Driver, cfg: getConfig(c.Level), jsonOutput: c.Json, diff --git a/cmd/dbc/uninstall_test.go b/cmd/dbc/uninstall_test.go index 20a6bd70..f401d763 100644 --- a/cmd/dbc/uninstall_test.go +++ b/cmd/dbc/uninstall_test.go @@ -84,13 +84,13 @@ func (suite *SubcommandTestSuite) TestUninstallMultipleLocations() { // Install to Env first m := InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.FileExists(filepath.Join(suite.tempdir, "test-driver-1.toml")) // Then System (here, we fake it as $tempdir/etc/adbc) m = InstallCmd{Driver: "test-driver-1", Level: config.ConfigSystem}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) installModel := m.(progressiveInstallModel) installModel.cfg.Location = filepath.Join(suite.tempdir, "root", installModel.cfg.Location) m = installModel // <- We need to reassign to make the change stick @@ -99,7 +99,7 @@ func (suite *SubcommandTestSuite) TestUninstallMultipleLocations() { // Uninstall from Env level m = UninstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.NoFileExists(filepath.Join(suite.tempdir, "test-driver-1.toml")) @@ -113,20 +113,20 @@ func (suite *SubcommandTestSuite) TestUninstallDriverTwice() { // Install to Env first m := InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.FileExists(filepath.Join(suite.tempdir, "test-driver-1.toml")) // Uninstall from Env level m = UninstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.NoFileExists(filepath.Join(suite.tempdir, "test-driver-1.toml")) // Uninstall from Env level m = UninstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r ", "\nError: failed to find driver `test-driver-1` in order to uninstall it: searched "+suite.tempdir, suite.runCmdErr(m)) } @@ -139,13 +139,13 @@ func (suite *SubcommandTestSuite) TestUninstallMultipleLocationsNonDefault() { // Install to Env first m := InstallCmd{Driver: "test-driver-1", Level: config.ConfigEnv}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.FileExists(filepath.Join(suite.tempdir, "test-driver-1.toml")) // Then System (here, we fake it as $tempdir/etc/adbc) m = InstallCmd{Driver: "test-driver-1", Level: config.ConfigSystem}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) installModel := m.(progressiveInstallModel) installModel.cfg.Location = filepath.Join(suite.tempdir, "root", installModel.cfg.Location) m = installModel // <- We need to reassign to make the change stick @@ -154,7 +154,7 @@ func (suite *SubcommandTestSuite) TestUninstallMultipleLocationsNonDefault() { // Then uninstall System (again, faked as $tempdir/etc/adbc) m = UninstallCmd{Driver: "test-driver-1", Level: config.ConfigSystem}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) uninstallModel := m.(uninstallModel) uninstallModel.cfg.Location = filepath.Join(suite.tempdir, "root", uninstallModel.cfg.Location) m = uninstallModel // <- We need to reassign to make the change stick @@ -166,7 +166,7 @@ func (suite *SubcommandTestSuite) TestUninstallMultipleLocationsNonDefault() { func (suite *SubcommandTestSuite) TestUninstallManifestOnlyDriver() { m := InstallCmd{Driver: "test-driver-manifest-only", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r[✓] searching\r\n[✓] downloading\r\n[✓] installing\r\n[✓] verifying signature\r\n", "\nInstalled test-driver-manifest-only 1.0.0 to "+suite.Dir()+ @@ -183,7 +183,7 @@ func (suite *SubcommandTestSuite) TestUninstallManifestOnlyDriver() { // Now uninstall and verify we clean up m = UninstallCmd{Driver: "test-driver-manifest-only", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.validateOutput("\r ", "Driver `test-driver-manifest-only` uninstalled successfully!", suite.runCmd(m)) suite.driverIsNotInstalled("test-driver-manifest-only") suite.NoDirExists(filepath.Join(suite.Dir(), new_sidecar_path)) @@ -196,7 +196,7 @@ func (suite *SubcommandTestSuite) TestUninstallInvalidManifest() { } m := InstallCmd{Driver: "test-driver-invalid-manifest", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) suite.runCmd(m) suite.FileExists(filepath.Join(suite.Dir(), "test-driver-invalid-manifest.toml")) @@ -243,7 +243,7 @@ func (suite *SubcommandTestSuite) TestUninstallRemovesSymlink() { // Install a driver m := InstallCmd{Driver: "test-driver-1", Level: suite.configLevel}. - GetModelCustom(baseModel{getDriverRegistry: getTestDriverRegistry, downloadPkg: downloadTestPkg}) + GetModelCustom(testBaseModel()) _ = suite.runCmd(m) suite.driverIsInstalled("test-driver-1", true) diff --git a/config/current.go b/cmd/dbc/view_config_tui.go similarity index 83% rename from config/current.go rename to cmd/dbc/view_config_tui.go index a91f0129..10d5a5f9 100644 --- a/config/current.go +++ b/cmd/dbc/view_config_tui.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package config +package main import ( "strings" @@ -20,7 +20,7 @@ import ( "charm.land/bubbles/v2/list" tea "charm.land/bubbletea/v2" "charm.land/lipgloss/v2" - "github.com/columnar-tech/dbc" + "github.com/columnar-tech/dbc/config" ) var ( @@ -35,10 +35,10 @@ var ( BorderStyle(lipgloss.HiddenBorder()) ) -type Model struct { +type configViewModel struct { Prev tea.Model - Drivers []DriverInfo + Drivers []config.DriverInfo list list.Model } @@ -46,7 +46,7 @@ var ( configStyle = lipgloss.NewStyle().PaddingLeft(2).Foreground(lipgloss.Color("170")) ) -type driverItem DriverInfo +type driverItem config.DriverInfo func (d driverItem) FilterValue() string { return d.ID } func (d driverItem) String() string { @@ -65,7 +65,7 @@ func (d driverItem) View() string { return configStyle.Render(sb.String()) } -func toListItems(drivers []DriverInfo) []list.Item { +func toListItems(drivers []config.DriverInfo) []list.Item { items := make([]list.Item, len(drivers)) for i, d := range drivers { items[i] = driverItem(d) @@ -73,17 +73,17 @@ func toListItems(drivers []DriverInfo) []list.Item { return items } -func (m Model) Init() tea.Cmd { +func (m configViewModel) Init() tea.Cmd { return func() tea.Msg { - return append(FindDriverConfigs(ConfigUser), - FindDriverConfigs(ConfigSystem)...) + return append(config.FindDriverConfigs(config.ConfigUser), + config.FindDriverConfigs(config.ConfigSystem)...) } } -func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { +func (m configViewModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg := msg.(type) { - case []DriverInfo: - m.list = list.New(toListItems(msg), dbc.SimpleItemDelegate{Prompt: ">"}, 20, 14) + case []config.DriverInfo: + m.list = list.New(toListItems(msg), SimpleItemDelegate{Prompt: ">"}, 20, 14) m.list.Title = "Available Drivers" m.list.SetShowStatusBar(false) m.list.SetFilteringEnabled(false) @@ -107,7 +107,7 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, cmd } -func (m Model) View() tea.View { +func (m configViewModel) View() tea.View { var sb strings.Builder sb.WriteString("DBC Driver Config\n\n") // sb.WriteString(configStyle.Render("System Driver Directory: "+systemDriversDir, diff --git a/config/config.go b/config/config.go index 071ea193..875662da 100644 --- a/config/config.go +++ b/config/config.go @@ -183,6 +183,7 @@ func loadConfig(lvl ConfigLevel) Config { maps.Copy(finalDrivers, drivers) } cfg.Exists, cfg.Drivers = len(finalDrivers) > 0, finalDrivers + return cfg } drivers, err := loadDir(cfg.Location) @@ -249,7 +250,9 @@ func InflateTarball(f *os.File, outDir string) (Manifest, error) { defer f.Close() var m Manifest - f.Seek(0, io.SeekStart) + if _, err := f.Seek(0, io.SeekStart); err != nil { + return m, fmt.Errorf("could not seek to start: %w", err) + } rdr, err := gzip.NewReader(f) if err != nil { return m, fmt.Errorf("could not create gzip reader: %w", err) @@ -423,11 +426,11 @@ func UninstallDriverShared(info DriverInfo) error { // Driver.shared is not a valid path (it's just a name), so this trick doesn't // work. We do want to clean this folder up so here we guess what it is and // try to remove it e.g., "somedriver_macos_arm64_v1.2.3." - extra_folder := fmt.Sprintf("%s_%s_v%s", info.ID, platformTuple, info.Version) - extra_folder = filepath.Clean(extra_folder) - finfo, err := root.Stat(extra_folder) - if err == nil && finfo.IsDir() && extra_folder != "." { - _ = root.RemoveAll(extra_folder) + extraFolder := fmt.Sprintf("%s_%s_v%s", info.ID, platformTuple, info.Version) + extraFolder = filepath.Clean(extraFolder) + finfo, err := root.Stat(extraFolder) + if err == nil && finfo.IsDir() && extraFolder != "." { + _ = root.RemoveAll(extraFolder) // ignore errors } diff --git a/config/config_api_test.go b/config/config_api_test.go new file mode 100644 index 00000000..8422caaf --- /dev/null +++ b/config/config_api_test.go @@ -0,0 +1,370 @@ +// Copyright 2026 Columnar Technologies 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. + +//go:build !windows + +package config_test + +import ( + "archive/tar" + "compress/gzip" + "io" + "os" + "path/filepath" + "testing" + + "github.com/Masterminds/semver/v3" + "github.com/columnar-tech/dbc/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testManifestTOML = ` +name = 'Test Driver' +publisher = 'Test Publisher' +license = 'MIT' +version = '1.2.3' +source = 'dbc' + +[ADBC] +version = '1.1.0' + +[Driver] +entrypoint = 'AdbcDriverInit' + +[Driver.shared] +linux_amd64 = '/path/to/driver.so' +` + +func makeTestDriverInfo(id string, filePath string) config.DriverInfo { + di := config.DriverInfo{ + ID: id, + FilePath: filePath, + Name: "Test Driver", + Publisher: "Test Publisher", + License: "MIT", + Version: semver.MustParse("1.2.3"), + Source: "dbc", + } + di.Driver.Entrypoint = "AdbcDriverInit" + di.Driver.Shared.Set("linux_amd64", filepath.Join(filePath, id, "driver.so")) + return di +} + +func TestGet(t *testing.T) { + t.Run("returns_all_levels", func(t *testing.T) { + configs := config.Get() + require.NotNil(t, configs) + assert.Contains(t, configs, config.ConfigSystem) + assert.Contains(t, configs, config.ConfigEnv) + }) + + t.Run("env_config_from_adbc_driver_path", func(t *testing.T) { + tmpDir := t.TempDir() + manifestPath := filepath.Join(tmpDir, "mydriver.toml") + require.NoError(t, os.WriteFile(manifestPath, []byte(testManifestTOML), 0644)) + + t.Setenv("ADBC_DRIVER_PATH", tmpDir) + + configs := config.Get() + envCfg := configs[config.ConfigEnv] + assert.True(t, envCfg.Exists) + assert.Contains(t, envCfg.Drivers, "mydriver") + }) + + t.Run("empty_env_path_yields_empty_config", func(t *testing.T) { + t.Setenv("ADBC_DRIVER_PATH", "") + t.Setenv("VIRTUAL_ENV", "") + t.Setenv("CONDA_PREFIX", "") + + configs := config.Get() + envCfg := configs[config.ConfigEnv] + assert.Equal(t, "", envCfg.Location) + assert.False(t, envCfg.Exists) + }) +} + +func TestInflateTarball(t *testing.T) { + t.Run("valid_tarball", func(t *testing.T) { + f, err := os.Open(filepath.Join("..", "cmd", "dbc", "testdata", "test-driver-1.tar.gz")) + require.NoError(t, err) + + outDir := t.TempDir() + manifest, err := config.InflateTarball(f, outDir) + require.NoError(t, err) + + assert.Equal(t, "Test Driver 1", manifest.Name) + assert.Equal(t, "1.0.0", manifest.Version.String()) + assert.NotEmpty(t, manifest.Files.Driver) + assert.FileExists(t, filepath.Join(outDir, manifest.Files.Driver)) + assert.FileExists(t, filepath.Join(outDir, manifest.Files.Signature)) + }) + + t.Run("invalid_gzip", func(t *testing.T) { + f, err := os.CreateTemp(t.TempDir(), "bad-*.tar.gz") + require.NoError(t, err) + _, _ = f.Write([]byte("this is not valid gzip data")) + _, _ = f.Seek(0, io.SeekStart) + + _, err = config.InflateTarball(f, t.TempDir()) + assert.Error(t, err) + }) + + t.Run("directory_entry_rejected", func(t *testing.T) { + f, err := os.CreateTemp(t.TempDir(), "dir-*.tar.gz") + require.NoError(t, err) + + gw := gzip.NewWriter(f) + tw := tar.NewWriter(gw) + require.NoError(t, tw.WriteHeader(&tar.Header{ + Name: "subdir/", + Typeflag: tar.TypeDir, + Mode: 0755, + })) + require.NoError(t, tw.Close()) + require.NoError(t, gw.Close()) + _, _ = f.Seek(0, io.SeekStart) + + _, err = config.InflateTarball(f, t.TempDir()) + assert.ErrorContains(t, err, "directory entry") + }) +} + +func TestInstallDriver(t *testing.T) { + t.Run("success", func(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("ADBC_DRIVER_PATH", tmpDir) + + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + f, err := os.Open(filepath.Join("..", "cmd", "dbc", "testdata", "test-driver-1.tar.gz")) + require.NoError(t, err) + + manifest, err := config.InstallDriver(cfg, "test-driver-1", f) + require.NoError(t, err) + + assert.Equal(t, "test-driver-1", manifest.DriverInfo.ID) + assert.Equal(t, "dbc", manifest.DriverInfo.Source) + assert.Equal(t, "Test Driver 1", manifest.DriverInfo.Name) + assert.Equal(t, "1.0.0", manifest.DriverInfo.Version.String()) + + platformTuple := config.PlatformTuple() + sharedPath := manifest.DriverInfo.Driver.Shared.Get(platformTuple) + assert.NotEmpty(t, sharedPath) + assert.FileExists(t, sharedPath) + }) + + t.Run("invalid_tarball", func(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + f, err := os.CreateTemp(t.TempDir(), "bad-*.tar.gz") + require.NoError(t, err) + _, _ = f.Write([]byte("not a tarball")) + _, _ = f.Seek(0, io.SeekStart) + + _, err = config.InstallDriver(cfg, "bad-driver", f) + assert.Error(t, err) + }) +} + +func TestGetDriver(t *testing.T) { + t.Run("found_in_env_config", func(t *testing.T) { + tmpDir := t.TempDir() + require.NoError(t, os.WriteFile( + filepath.Join(tmpDir, "mydriver.toml"), + []byte(testManifestTOML), + 0644, + )) + + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + di, err := config.GetDriver(cfg, "mydriver") + require.NoError(t, err) + assert.Equal(t, "mydriver", di.ID) + assert.Equal(t, "Test Driver", di.Name) + assert.Equal(t, "1.2.3", di.Version.String()) + }) + + t.Run("not_found", func(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + _, err := config.GetDriver(cfg, "nonexistent") + assert.Error(t, err) + }) + + t.Run("multi_path_env_finds_first_match", func(t *testing.T) { + dir1 := t.TempDir() + dir2 := t.TempDir() + require.NoError(t, os.WriteFile( + filepath.Join(dir2, "driver2.toml"), + []byte(testManifestTOML), + 0644, + )) + + cfg := config.Config{ + Level: config.ConfigEnv, + Location: dir1 + string(filepath.ListSeparator) + dir2, + } + + di, err := config.GetDriver(cfg, "driver2") + require.NoError(t, err) + assert.Equal(t, "driver2", di.ID) + }) +} + +func TestCreateManifest(t *testing.T) { + t.Run("success", func(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.Config{ + Level: config.ConfigEnv, + Location: tmpDir, + } + + di := makeTestDriverInfo("mydriver", tmpDir) + + err := config.CreateManifest(cfg, di) + require.NoError(t, err) + + assert.FileExists(t, filepath.Join(tmpDir, "mydriver.toml")) + }) + + t.Run("creates_location_if_absent", func(t *testing.T) { + newDir := filepath.Join(t.TempDir(), "nonexistent", "subdir") + cfg := config.Config{ + Level: config.ConfigEnv, + Location: newDir, + } + + di := makeTestDriverInfo("newdriver", newDir) + + err := config.CreateManifest(cfg, di) + require.NoError(t, err) + + assert.FileExists(t, filepath.Join(newDir, "newdriver.toml")) + }) +} + +func TestFindDriverConfigs(t *testing.T) { + t.Run("empty_location_returns_empty", func(t *testing.T) { + t.Setenv("ADBC_DRIVER_PATH", "") + t.Setenv("VIRTUAL_ENV", "") + t.Setenv("CONDA_PREFIX", "") + + drivers := config.FindDriverConfigs(config.ConfigEnv) + assert.Empty(t, drivers) + }) + + t.Run("returns_installed_drivers", func(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("ADBC_DRIVER_PATH", tmpDir) + + require.NoError(t, os.WriteFile( + filepath.Join(tmpDir, "driver1.toml"), + []byte(testManifestTOML), + 0644, + )) + require.NoError(t, os.WriteFile( + filepath.Join(tmpDir, "driver2.toml"), + []byte(testManifestTOML), + 0644, + )) + + drivers := config.FindDriverConfigs(config.ConfigEnv) + require.Len(t, drivers, 2) + + ids := make([]string, len(drivers)) + for i, d := range drivers { + ids[i] = d.ID + } + assert.ElementsMatch(t, []string{"driver1", "driver2"}, ids) + }) + + t.Run("nonexistent_path_returns_empty", func(t *testing.T) { + t.Setenv("ADBC_DRIVER_PATH", "/nonexistent/path/that/does/not/exist") + + drivers := config.FindDriverConfigs(config.ConfigEnv) + assert.Empty(t, drivers) + }) +} + +func TestUninstallDriverShared(t *testing.T) { + t.Run("dbc_source_removes_driver_dir", func(t *testing.T) { + tmpDir := t.TempDir() + + driverDir := filepath.Join(tmpDir, "test-driver-1_linux_amd64_v1.0.0") + require.NoError(t, os.MkdirAll(driverDir, 0755)) + + driverPath := filepath.Join(driverDir, "driver.so") + require.NoError(t, os.WriteFile(driverPath, []byte("fake so"), 0644)) + + di := config.DriverInfo{ + ID: "test-driver-1", + FilePath: tmpDir, + Source: "dbc", + } + di.Driver.Shared.Set("linux_amd64", driverPath) + + err := config.UninstallDriverShared(di) + require.NoError(t, err) + + assert.NoDirExists(t, driverDir) + }) + + t.Run("non_dbc_source_removes_file_only", func(t *testing.T) { + tmpDir := t.TempDir() + + driverPath := filepath.Join(tmpDir, "driver.so") + require.NoError(t, os.WriteFile(driverPath, []byte("fake so"), 0644)) + + di := config.DriverInfo{ + ID: "external-driver", + FilePath: tmpDir, + Source: "external", + } + di.Driver.Shared.Set("linux_amd64", driverPath) + + err := config.UninstallDriverShared(di) + require.NoError(t, err) + + assert.NoFileExists(t, driverPath) + }) + + t.Run("missing_file_is_tolerated_for_non_dbc", func(t *testing.T) { + tmpDir := t.TempDir() + + di := config.DriverInfo{ + ID: "missing-driver", + FilePath: tmpDir, + Source: "external", + } + di.Driver.Shared.Set("linux_amd64", filepath.Join(tmpDir, "nonexistent.so")) + + err := config.UninstallDriverShared(di) + assert.NoError(t, err) + }) +} diff --git a/config/driver.go b/config/driver.go index f74b74b8..c53b7767 100644 --- a/config/driver.go +++ b/config/driver.go @@ -24,7 +24,6 @@ import ( "strings" "github.com/Masterminds/semver/v3" - "github.com/columnar-tech/dbc" "github.com/pelletier/go-toml/v2" ) @@ -43,17 +42,6 @@ type Manifest struct { } `toml:"PostInstall,omitempty"` } -func (m Manifest) ToPackageInfo() dbc.PkgInfo { - return dbc.PkgInfo{ - Driver: dbc.Driver{ - Title: m.Name, - Path: m.ID, - License: m.License, - }, - Version: m.Version, - } -} - type DriverInfo struct { ID string FilePath string @@ -205,8 +193,8 @@ func createDriverManifest(location string, driver DriverInfo) error { } } - manifest_path := filepath.Join(location, driver.ID+".toml") - f, err := os.Create(manifest_path) + manifestPath := filepath.Join(location, driver.ID+".toml") + f, err := os.Create(manifestPath) if err != nil { return fmt.Errorf("error creating manifest %s: %w", driver.ID, err) } @@ -220,7 +208,7 @@ func createDriverManifest(location string, driver DriverInfo) error { // installing. // // TODO: Remove this when the driver managers are fixed (>=1.8.1). - createManifestSymlink(location, driver.ID, manifest_path) + createManifestSymlink(location, driver.ID, manifestPath) toEncode := tomlDriverInfo{ ManifestVersion: currentManifestVersion, diff --git a/dbc_api_test.go b/dbc_api_test.go new file mode 100644 index 00000000..472f5b08 --- /dev/null +++ b/dbc_api_test.go @@ -0,0 +1,462 @@ +// Copyright 2026 Columnar Technologies 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 dbc_test + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/Masterminds/semver/v3" + "github.com/columnar-tech/dbc" + "github.com/go-faster/yaml" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var testServer *httptest.Server + +type testTransport struct{} + +func (t *testTransport) RoundTrip(req *http.Request) (*http.Response, error) { + srvURL, _ := url.Parse(testServer.URL) + newURL := *req.URL + newURL.Scheme = srvURL.Scheme + newURL.Host = srvURL.Host + newReq := req.Clone(req.Context()) + newReq.URL = &newURL + newReq.Host = req.URL.Host + return http.DefaultTransport.RoundTrip(newReq) +} + +func TestMain(m *testing.M) { + indexData, err := os.ReadFile(filepath.Join("cmd", "dbc", "testdata", "test_index.yaml")) + if err != nil { + panic("cannot read test_index.yaml: " + err.Error()) + } + + tarballData, err := os.ReadFile(filepath.Join("cmd", "dbc", "testdata", "test-driver-1.tar.gz")) + if err != nil { + panic("cannot read test-driver-1.tar.gz: " + err.Error()) + } + + testServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/index.yaml": + w.Header().Set("Content-Type", "application/yaml") + w.Write(indexData) + case "/test-driver-1.tar.gz": + w.Header().Set("Content-Type", "application/gzip") + w.Header().Set("Content-Length", fmt.Sprint(len(tarballData))) + w.Write(tarballData) + default: + http.NotFound(w, r) + } + })) + + origClient := dbc.DefaultClient + dbc.DefaultClient = &http.Client{Transport: &testTransport{}} + + code := m.Run() + + testServer.Close() + dbc.DefaultClient = origClient + os.Exit(code) +} + +func mustParseURL(s string) *url.URL { + u, err := url.Parse(s) + if err != nil { + panic(err) + } + return u +} + +func loadTestDrivers(t *testing.T, registry *dbc.Registry) []dbc.Driver { + t.Helper() + + f, err := os.Open(filepath.Join("cmd", "dbc", "testdata", "test_index.yaml")) + require.NoError(t, err) + defer f.Close() + + var index struct { + Drivers []dbc.Driver `yaml:"drivers"` + } + require.NoError(t, yaml.NewDecoder(f).Decode(&index)) + for i := range index.Drivers { + index.Drivers[i].Registry = registry + } + return index.Drivers +} + +func findDriver(t *testing.T, drivers []dbc.Driver, path string) dbc.Driver { + t.Helper() + for _, d := range drivers { + if d.Path == path { + return d + } + } + t.Fatalf("driver %q not found in list", path) + return dbc.Driver{} +} + +func extractTarFiles(t *testing.T, tarPath string, names ...string) map[string][]byte { + t.Helper() + + f, err := os.Open(tarPath) + require.NoError(t, err) + defer f.Close() + + gr, err := gzip.NewReader(f) + require.NoError(t, err) + defer gr.Close() + + result := make(map[string][]byte) + tr := tar.NewReader(gr) + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + for _, name := range names { + if hdr.Name == name { + data, err := io.ReadAll(tr) + require.NoError(t, err) + result[name] = data + } + } + } + return result +} + +func TestGetDriverList(t *testing.T) { + t.Run("success", func(t *testing.T) { + drivers, _ := dbc.GetDriverList() + require.NotEmpty(t, drivers, "GetDriverList must return at least one driver") + + var paths []string + for _, d := range drivers { + paths = append(paths, d.Path) + } + assert.Contains(t, paths, "test-driver-1") + assert.Contains(t, paths, "test-driver-2") + assert.Contains(t, paths, "test-driver-only-pre") + + driver1 := findDriver(t, drivers, "test-driver-1") + assert.Equal(t, "Test Driver 1", driver1.Title) + assert.Equal(t, "MIT", driver1.License) + assert.NotNil(t, driver1.Registry) + assert.NotEmpty(t, driver1.Registry.BaseURL) + }) +} + +func TestDriverHasNonPrerelease(t *testing.T) { + registry := &dbc.Registry{BaseURL: mustParseURL("https://registry.example.com")} + drivers := loadTestDrivers(t, registry) + + t.Run("has_stable_releases", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + assert.True(t, d.HasNonPrerelease()) + }) + + t.Run("only_prerelease", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-only-pre") + assert.False(t, d.HasNonPrerelease()) + }) + + t.Run("empty_pkginfo", func(t *testing.T) { + d := dbc.Driver{} + assert.False(t, d.HasNonPrerelease()) + }) +} + +func TestDriverVersions(t *testing.T) { + registry := &dbc.Registry{BaseURL: mustParseURL("https://registry.example.com")} + drivers := loadTestDrivers(t, registry) + + t.Run("linux_amd64", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + versions := d.Versions("linux_amd64") + require.Len(t, versions, 2) + assert.Equal(t, "1.0.0", versions[0].String()) + assert.Equal(t, "1.1.0", versions[1].String()) + }) + + t.Run("unknown_platform", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + versions := d.Versions("unknown_platform") + assert.Empty(t, versions) + }) + + t.Run("prerelease_included", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-2") + versions := d.Versions("linux_amd64") + require.Len(t, versions, 3) + assert.Equal(t, "2.0.0", versions[0].String()) + assert.Equal(t, "2.1.0-beta.1", versions[1].String()) + assert.Equal(t, "2.1.0", versions[2].String()) + }) +} + +func TestDriverMaxVersion(t *testing.T) { + registry := &dbc.Registry{BaseURL: mustParseURL("https://registry.example.com")} + drivers := loadTestDrivers(t, registry) + + t.Run("returns_highest_version", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + max, ok := d.MaxVersion() + require.True(t, ok) + assert.Equal(t, "1.1.0", max.Version.String()) + }) + + t.Run("prerelease_can_be_max", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-only-pre") + max, ok := d.MaxVersion() + require.True(t, ok) + assert.Equal(t, "0.9.0-alpha.1", max.Version.String()) + }) +} + +func TestDriverGetPackage(t *testing.T) { + registry := &dbc.Registry{BaseURL: mustParseURL("https://registry.example.com")} + drivers := loadTestDrivers(t, registry) + + t.Run("latest_version_linux_amd64", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + pkg, err := d.GetPackage(nil, "linux_amd64", false) + require.NoError(t, err) + assert.Equal(t, "1.1.0", pkg.Version.String()) + assert.Equal(t, "linux_amd64", pkg.PlatformTuple) + assert.NotNil(t, pkg.Path) + }) + + t.Run("specific_version", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + v := semver.MustParse("1.0.0") + pkg, err := d.GetPackage(v, "linux_amd64", false) + require.NoError(t, err) + assert.Equal(t, "1.0.0", pkg.Version.String()) + }) + + t.Run("version_not_found", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + v := semver.MustParse("9.9.9") + _, err := d.GetPackage(v, "linux_amd64", false) + assert.Error(t, err) + assert.Contains(t, err.Error(), "version 9.9.9 not found") + }) + + t.Run("prerelease_filtered_out", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-2") + pkg, err := d.GetPackage(nil, "linux_amd64", false) + require.NoError(t, err) + assert.Equal(t, "2.1.0", pkg.Version.String()) + assert.Equal(t, "", pkg.Version.Prerelease()) + }) + + t.Run("allow_prerelease", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-2") + pkg, err := d.GetPackage(nil, "linux_amd64", true) + require.NoError(t, err) + assert.Equal(t, "2.1.0", pkg.Version.String()) + }) + + t.Run("only_prerelease_not_allowed", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-only-pre") + _, err := d.GetPackage(nil, "linux_amd64", false) + require.Error(t, err) + assert.Contains(t, err.Error(), "prerelease versions filtered out") + }) + + t.Run("no_packages_for_platform", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + _, err := d.GetPackage(nil, "nonexistent_platform", false) + assert.Error(t, err) + }) + + t.Run("empty_driver", func(t *testing.T) { + d := dbc.Driver{Path: "empty"} + _, err := d.GetPackage(nil, "linux_amd64", false) + assert.Error(t, err) + assert.Contains(t, err.Error(), "driver `empty` not found") + }) +} + +func TestDriverGetWithConstraint(t *testing.T) { + registry := &dbc.Registry{BaseURL: mustParseURL("https://registry.example.com")} + drivers := loadTestDrivers(t, registry) + + t.Run("matches_constraint", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + c, err := semver.NewConstraint(">=1.0.0 <1.1.0") + require.NoError(t, err) + + pkg, err := d.GetWithConstraint(c, "linux_amd64") + require.NoError(t, err) + assert.Equal(t, "1.0.0", pkg.Version.String()) + }) + + t.Run("returns_highest_matching", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + c, err := semver.NewConstraint(">=1.0.0") + require.NoError(t, err) + + pkg, err := d.GetWithConstraint(c, "linux_amd64") + require.NoError(t, err) + assert.Equal(t, "1.1.0", pkg.Version.String()) + }) + + t.Run("no_matching_version", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + c, err := semver.NewConstraint(">=99.0.0") + require.NoError(t, err) + + _, err = d.GetWithConstraint(c, "linux_amd64") + assert.Error(t, err) + assert.Contains(t, err.Error(), "no package found for driver") + }) + + t.Run("no_matching_platform", func(t *testing.T) { + d := findDriver(t, drivers, "test-driver-1") + c, err := semver.NewConstraint(">=1.0.0") + require.NoError(t, err) + + _, err = d.GetWithConstraint(c, "nonexistent_platform") + assert.Error(t, err) + }) + + t.Run("empty_pkginfo", func(t *testing.T) { + d := dbc.Driver{Path: "empty"} + c, err := semver.NewConstraint(">=1.0.0") + require.NoError(t, err) + + _, err = d.GetWithConstraint(c, "linux_amd64") + assert.Error(t, err) + assert.Contains(t, err.Error(), "no package info available") + }) +} + +func TestPkgInfoDownloadPackage(t *testing.T) { + t.Run("no_url", func(t *testing.T) { + pkg := dbc.PkgInfo{ + Driver: dbc.Driver{Title: "test-driver"}, + Version: semver.MustParse("1.0.0"), + } + _, err := pkg.DownloadPackage(nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "no url set") + }) + + t.Run("success", func(t *testing.T) { + u := mustParseURL(testServer.URL + "/test-driver-1.tar.gz") + pkg := dbc.PkgInfo{ + Driver: dbc.Driver{Title: "test-driver-1"}, + Version: semver.MustParse("1.0.0"), + PlatformTuple: "linux_amd64", + Path: u, + } + + f, err := pkg.DownloadPackage(nil) + require.NoError(t, err) + require.NotNil(t, f) + defer os.Remove(f.Name()) + defer f.Close() + + fi, err := f.Stat() + require.NoError(t, err) + assert.Greater(t, fi.Size(), int64(0)) + }) + + t.Run("success_with_progress_callback", func(t *testing.T) { + u := mustParseURL(testServer.URL + "/test-driver-1.tar.gz") + pkg := dbc.PkgInfo{ + Driver: dbc.Driver{Title: "test-driver-1"}, + Version: semver.MustParse("1.0.0"), + PlatformTuple: "linux_amd64", + Path: u, + } + + var lastWritten, lastTotal int64 + f, err := pkg.DownloadPackage(func(written, total int64) { + lastWritten = written + lastTotal = total + }) + require.NoError(t, err) + require.NotNil(t, f) + defer os.Remove(f.Name()) + defer f.Close() + + assert.Greater(t, lastWritten, int64(0)) + _ = lastTotal + }) + + t.Run("server_error", func(t *testing.T) { + u := mustParseURL(testServer.URL + "/does-not-exist.tar.gz") + pkg := dbc.PkgInfo{ + Driver: dbc.Driver{Title: "test-driver-1"}, + Version: semver.MustParse("1.0.0"), + PlatformTuple: "linux_amd64", + Path: u, + } + + _, err := pkg.DownloadPackage(nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to download driver") + }) +} + +func TestSignedByColumnar(t *testing.T) { + t.Run("valid_signature", func(t *testing.T) { + tarPath := filepath.Join("cmd", "dbc", "testdata", "test-driver-1.tar.gz") + files := extractTarFiles(t, tarPath, + "test-driver-1-not-valid.so", + "test-driver-1-not-valid.so.sig", + ) + require.Contains(t, files, "test-driver-1-not-valid.so") + require.Contains(t, files, "test-driver-1-not-valid.so.sig") + + lib := bytes.NewReader(files["test-driver-1-not-valid.so"]) + sig := bytes.NewReader(files["test-driver-1-not-valid.so.sig"]) + + err := dbc.SignedByColumnar(lib, sig) + assert.NoError(t, err) + }) + + t.Run("invalid_signature", func(t *testing.T) { + lib := strings.NewReader("this is not a real driver binary") + sig := strings.NewReader("this is not a real pgp signature") + + err := dbc.SignedByColumnar(lib, sig) + assert.Error(t, err) + }) + + t.Run("empty_inputs", func(t *testing.T) { + lib := strings.NewReader("") + sig := strings.NewReader("") + + err := dbc.SignedByColumnar(lib, sig) + assert.Error(t, err) + }) +} diff --git a/drivers.go b/drivers.go index 231e0550..d5c172ed 100644 --- a/drivers.go +++ b/drivers.go @@ -15,6 +15,7 @@ package dbc import ( + "context" _ "embed" "errors" "fmt" @@ -36,7 +37,6 @@ import ( "github.com/ProtonMail/gopenpgp/v3/crypto" "github.com/columnar-tech/dbc/auth" "github.com/columnar-tech/dbc/internal" - "github.com/go-faster/yaml" "github.com/google/uuid" machineid "github.com/zeroshade/machine-id" ) @@ -61,17 +61,19 @@ func mustParseURL(u string) *url.URL { } var ( - registries = []Registry{ - {BaseURL: mustParseURL("https://dbc-cdn.columnar.tech")}, - {BaseURL: mustParseURL("https://" + auth.DefaultOauthURI())}, - } Version = "unknown" mid string uid uuid.UUID - // use this default client for all requests, - // it will add the dbc user-agent to all requests + // DefaultClient is the HTTP client used for all requests. + // + // Deprecated: Use NewClient with WithHTTPClient instead. + // DefaultClient must be set during program initialization, + // before any concurrent calls to GetDriverList or makereq. DefaultClient = http.DefaultClient + + setupOnce sync.Once + internalClient *http.Client ) type uaRoundTripper struct { @@ -81,6 +83,7 @@ type uaRoundTripper struct { // custom RoundTripper that sets the User-Agent header on any requests func (u *uaRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + req = req.Clone(req.Context()) req.Header.Set("User-Agent", u.userAgent) return u.RoundTripper.RoundTrip(req) } @@ -90,62 +93,65 @@ func init() { if ok && Version == "unknown" { Version = info.Main.Version } +} - if val := os.Getenv("DBC_BASE_URL"); val != "" { - registries = []Registry{ - {BaseURL: mustParseURL(val)}, - } - } - - userAgent := fmt.Sprintf("dbc-cli/%s (%s; %s)", - Version, runtime.GOOS, runtime.GOARCH) +func ensureSetup() { + setupOnce.Do(func() { + userAgent := fmt.Sprintf("dbc-cli/%s (%s; %s)", + Version, runtime.GOOS, runtime.GOARCH) - // many CI systems set CI=true in the env so let's check for that - if ci := os.Getenv("CI"); ci != "" { - if val, _ := strconv.ParseBool(ci); val { - userAgent += " CI" + if ci := os.Getenv("CI"); ci != "" { + if val, _ := strconv.ParseBool(ci); val { + userAgent += " CI" + } } - } - DefaultClient.Transport = &uaRoundTripper{ - RoundTripper: http.DefaultTransport, - userAgent: userAgent, - } + internalClient = &http.Client{ + Transport: &uaRoundTripper{ + RoundTripper: http.DefaultTransport, + userAgent: userAgent, + }, + } - mid, _ = machineid.ProtectedID() + mid, _ = machineid.ProtectedID() - // get user config dir - userdir, err := internal.GetUserConfigPath() - if err != nil { - // if we can't get the dir for some reason, just generate a new UUID - uid = uuid.New() - return - } - - fp := filepath.Join(userdir, "uid.uuid") - data, err := os.ReadFile(fp) - if err == nil { - if err = uid.UnmarshalBinary(data); err == nil { + userdir, err := internal.GetUserConfigPath() + if err != nil { + uid = uuid.New() return } - } - // if the file didn't exist or we couldn't parse it, generate a new uuid - // and then write a new file - uid = uuid.New() - // if we fail to create the dir or write the file, just ignore the error - // and use the fresh UUID - if err = os.MkdirAll(filepath.Dir(fp), 0o700); err == nil { - if data, err = uid.MarshalBinary(); err == nil { - os.WriteFile(fp, data, 0o600) + fp := filepath.Join(userdir, "uid.uuid") + data, err := os.ReadFile(fp) + if err == nil { + if err = uid.UnmarshalBinary(data); err == nil { + return + } } + + uid = uuid.New() + if err = os.MkdirAll(filepath.Dir(fp), 0o700); err == nil { + if data, err = uid.MarshalBinary(); err == nil { + os.WriteFile(fp, data, 0o600) + } + } + }) +} + +func getHTTPClient() *http.Client { + ensureSetup() + if DefaultClient != http.DefaultClient { + return DefaultClient } + return internalClient } func makereq(u string) (resp *http.Response, err error) { + ensureSetup() + uri, err := url.Parse(u) if err != nil { - return nil, fmt.Errorf("failed to parse URL %s: %w", uri, err) + return nil, fmt.Errorf("failed to parse URL %s: %w", u, err) } cred, err := auth.GetCredentials(uri) @@ -158,16 +164,22 @@ func makereq(u string) (resp *http.Response, err error) { q.Add("uid", uid.String()) uri.RawQuery = q.Encode() - req := http.Request{ - Method: http.MethodGet, - URL: uri, - Header: http.Header{}, - } - - if uri.Path == "/index.yaml" { - req.Header.Set("Accept", "application/yaml") + buildLegacyReq := func(token string) (*http.Request, error) { + urlCopy := *uri + r, err := http.NewRequestWithContext(context.Background(), http.MethodGet, urlCopy.String(), nil) + if err != nil { + return nil, err + } + if uri.Path == "/index.yaml" { + r.Header.Set("Accept", "application/yaml") + } + if token != "" { + r.Header.Set("Authorization", "Bearer "+token) + } + return r, nil } + token := "" if cred != nil { if auth.IsColumnarPrivateRegistry(uri) { // if we're accessing the private registry then attempt to @@ -176,23 +188,31 @@ func makereq(u string) (resp *http.Response, err error) { // trial or it is expired, then this will silently fail. _ = auth.FetchColumnarLicense(cred) } - req.Header.Set("Authorization", "Bearer "+cred.GetAuthToken()) + token = cred.GetAuthToken() } - resp, err = DefaultClient.Do(&req) + req, err := buildLegacyReq(token) + if err != nil { + return nil, fmt.Errorf("failed to build request: %w", err) + } + resp, err = getHTTPClient().Do(req) if err != nil { return } if resp.StatusCode == http.StatusUnauthorized && cred != nil { resp.Body.Close() - // Try refreshing the token - if !cred.Refresh() { - return nil, fmt.Errorf("failed to refresh auth token") + if err := cred.Refresh(); err != nil { + return nil, fmt.Errorf("failed to refresh auth token: %w", err) + } + retryReq, retryErr := buildLegacyReq(cred.GetAuthToken()) + if retryErr != nil { + return nil, fmt.Errorf("failed to build retry request: %w", retryErr) + } + resp, err = getHTTPClient().Do(retryReq) + if err != nil { + return } - - req.Header.Set("Authorization", "Bearer "+cred.GetAuthToken()) - resp, err = DefaultClient.Do(&req) } switch resp.StatusCode { @@ -208,65 +228,6 @@ func makereq(u string) (resp *http.Response, err error) { return resp, err } -func getDriverListFromIndex(index *Registry) ([]Driver, error) { - resp, err := makereq(index.BaseURL.JoinPath("/index.yaml").String()) - if err != nil { - return nil, fmt.Errorf("failed to fetch drivers: %w", err) - } - - if resp.StatusCode != http.StatusOK { - // ignore registries we aren't authorized to access - if resp.StatusCode == http.StatusUnauthorized { - return nil, nil - } - - return nil, fmt.Errorf("failed to fetch drivers: %s", resp.Status) - } - - defer resp.Body.Close() - drivers := struct { - Name string `yaml:"name"` - Drivers []Driver `yaml:"drivers"` - }{} - - err = yaml.NewDecoder(resp.Body).Decode(&drivers) - if err != nil { - return nil, fmt.Errorf("failed to parse driver registry index: %s", err) - } - - if drivers.Name != "" { - index.Name = drivers.Name - } - - // Set registry reference - for i := range drivers.Drivers { - drivers.Drivers[i].Registry = index - } - - result := drivers.Drivers - sort.Slice(result, func(i, j int) bool { - return result[i].Path < result[j].Path - }) - - return result, nil -} - -var getDrivers = sync.OnceValues(func() ([]Driver, error) { - var totalErr error - allDrivers := make([]Driver, 0) - for i := range registries { - drivers, err := getDriverListFromIndex(®istries[i]) - if err != nil { - totalErr = errors.Join(totalErr, fmt.Errorf("registry %s: %w", registries[i].BaseURL, err)) - continue - } - registries[i].Drivers = drivers - allDrivers = append(allDrivers, drivers...) - } - - return allDrivers, totalErr -}) - //go:embed columnar.pubkey var armoredPubKey string @@ -305,6 +266,7 @@ type PkgInfo struct { Path *url.URL } +// Deprecated: Use Client.Download instead. func (p PkgInfo) DownloadPackage(prog ProgressFunc) (*os.File, error) { if p.Path == nil { return nil, fmt.Errorf("cannot download package for %s: no url set", p.Driver.Title) @@ -317,6 +279,7 @@ func (p PkgInfo) DownloadPackage(prog ProgressFunc) (*os.File, error) { } if rsp.StatusCode != http.StatusOK { + rsp.Body.Close() return nil, fmt.Errorf("failed to download driver %s: %s", location, rsp.Status) } defer rsp.Body.Close() @@ -329,6 +292,7 @@ func (p PkgInfo) DownloadPackage(prog ProgressFunc) (*os.File, error) { output, err := os.Create(path.Join(tmpdir, fname)) if err != nil { + os.RemoveAll(tmpdir) return nil, fmt.Errorf("failed to create temp file to download to: %w", err) } @@ -341,6 +305,8 @@ func (p PkgInfo) DownloadPackage(prog ProgressFunc) (*os.File, error) { _, err = io.Copy(pw, rsp.Body) if err != nil { output.Close() + output = nil + os.RemoveAll(tmpdir) } return output, err } @@ -358,13 +324,20 @@ func (p pkginfo) GetPackage(d Driver, platformTuple string) (PkgInfo, error) { return PkgInfo{}, fmt.Errorf("no packages available for version %s", p.Version) } + if d.Registry == nil { + return PkgInfo{}, fmt.Errorf("cannot resolve package URL for %s: driver has no registry", d.Title) + } base := d.Registry.BaseURL for _, pkg := range p.Packages { if pkg.PlatformTuple == platformTuple { var uri *url.URL if pkg.URL != "" { - uri, _ = url.Parse(pkg.URL) + var err error + uri, err = url.Parse(pkg.URL) + if err != nil { + return PkgInfo{}, fmt.Errorf("invalid package URL %q: %w", pkg.URL, err) + } if !uri.IsAbs() { uri = base.JoinPath(pkg.URL) } @@ -403,7 +376,7 @@ type Driver struct { License string `yaml:"license"` Path string `yaml:"path"` URLs []string `yaml:"urls"` - DocsUrl string `yaml:"docs_url"` + DocsURL string `yaml:"docs_url"` PkgInfo []pkginfo `yaml:"pkginfo"` } @@ -434,7 +407,8 @@ func (d Driver) GetWithConstraint(c *semver.Constraints, platformTuple string) ( var result *pkginfo for pkg := range itr { if result == nil || pkg.Version.GreaterThan(result.Version) { - result = &pkg + found := pkg + result = &found } } @@ -462,8 +436,10 @@ func (d Driver) Versions(platformTuple string) semver.Collection { func (d Driver) GetPackage(version *semver.Version, platformTuple string, allowPrerelease bool) (PkgInfo, error) { pkglist := d.PkgInfo - // Filter out pre-releases and record whether any pre-releases were filtered - // out so we can produce a more helpful error message + // Filter prereleases when no specific stable version is requested. + // When version is a specific stable release (Prerelease() == ""), + // filtering is unnecessary — the exact-match search below will + // only match the requested stable version. if !allowPrerelease && (version == nil || version.Prerelease() != "") { hadPackages := len(d.PkgInfo) > 0 pkglist = slices.Collect(filter(slices.Values(d.PkgInfo), func(p pkginfo) bool { @@ -488,6 +464,9 @@ func (d Driver) GetPackage(version *semver.Version, platformTuple string, allowP return p.Version.Equal(version) }) if idx == -1 { + if !allowPrerelease && version.Prerelease() != "" { + return PkgInfo{}, fmt.Errorf("version %s is a prerelease; use --pre to allow it", version) + } return PkgInfo{}, fmt.Errorf("version %s not found", version) } pkg = pkglist[idx] @@ -496,16 +475,75 @@ func (d Driver) GetPackage(version *semver.Version, platformTuple string, allowP return pkg.GetPackage(d, platformTuple) } -func (d Driver) MaxVersion() pkginfo { - return slices.MaxFunc(d.PkgInfo, func(a, b pkginfo) int { +func (d Driver) MaxVersion() (VersionInfo, bool) { + if len(d.PkgInfo) == 0 { + return VersionInfo{}, false + } + p := slices.MaxFunc(d.PkgInfo, func(a, b pkginfo) int { return a.Version.Compare(b.Version) }) + pkgs := make([]PackageInfo, 0, len(p.Packages)) + for _, pkg := range p.Packages { + pkgs = append(pkgs, PackageInfo{ + Platform: pkg.PlatformTuple, + URL: pkg.URL, + }) + } + return VersionInfo{Version: p.Version, Packages: pkgs}, true +} + +// PackageInfo holds the platform and raw URL string for a single package entry. +// The URL may be relative (joined against the registry base URL) or absolute. +type PackageInfo struct { + Platform string + URL string +} + +// VersionInfo holds the version and its associated packages for a driver. +type VersionInfo struct { + Version *semver.Version + Packages []PackageInfo } +// AllVersions returns all version/package entries for the driver as exported +// VersionInfo values. This allows callers outside the dbc package to iterate +// over every version and platform without needing access to the unexported +// pkginfo type. +func (d Driver) AllVersions() []VersionInfo { + result := make([]VersionInfo, 0, len(d.PkgInfo)) + for _, pi := range d.PkgInfo { + pkgs := make([]PackageInfo, 0, len(pi.Packages)) + for _, p := range pi.Packages { + pkgs = append(pkgs, PackageInfo{ + Platform: p.PlatformTuple, + URL: p.URL, + }) + } + result = append(result, VersionInfo{ + Version: pi.Version, + Packages: pkgs, + }) + } + return result +} + +// GetDriverList returns a list of all available drivers from all configured registries. +// +// Deprecated: Use NewClient and Client.Search instead. func GetDriverList() ([]Driver, error) { - return getDrivers() + ensureSetup() + var opts []Option + if val := os.Getenv("DBC_BASE_URL"); val != "" { + opts = append(opts, WithBaseURL(val)) + } + c, err := NewClient(append(opts, WithHTTPClient(getHTTPClient()))...) + if err != nil { + return nil, err + } + return c.Search("") } +// Deprecated: Signature verification is now handled internally by Client.Install. // SignedByColumnar returns nil if the library was signed by // the columnar public key (embedded in the CLI) or an error // otherwise.