feat(auth): automatically fetch Columnar license on login (#223)

When using `dbc auth login` to login to the Columnar private registry,
we'll check the token and see if it contains the `trial_start` claim,
downloading the license if it exists.

There's some more work that needs to happen here so I'm leaving this
marked as WIP, such as:

* Picking a domain for the cloudflare workers
* making the worker domain configurable via env var
* unit tests
* verifying this workflow makes sense

---------

Co-authored-by: Bryce Mecum <petridish@gmail.com>

Author: zeroshade

auth/credentials.go

@@ -18,6 +18,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io/fs"
 	"net/http"
 	"net/url"
 	"os"
@@ -26,6 +27,7 @@ import (
 	"slices"
 	"sync"
 
+	"github.com/golang-jwt/jwt/v5"
 	"github.com/pelletier/go-toml/v2"
 )
 
@@ -265,3 +267,80 @@ func UpdateCreds() error {
 		Credentials: loadedCredentials,
 	})
 }
+
+func IsColumnarPrivateRegistry(u *url.URL) bool {
+	return u.Host == DefaultOauthURI
+}
+
+const licenseURI = "https://heimdall.columnar.tech/trial_license"
+
+var (
+	ErrNoTrialLicense = errors.New("no trial license found")
+	ErrTrialExpired   = errors.New("trial license has expired")
+)
+
+func FetchColumnarLicense(cred *Credential) error {
+	licensePath := filepath.Join(filepath.Dir(credPath), "columnar.lic")
+	_, err := os.Stat(licensePath)
+	if err == nil { // license exists already
+		return nil
+	}
+
+	if !errors.Is(err, fs.ErrNotExist) {
+		return err
+	}
+
+	var authToken string
+	switch cred.Type {
+	case TypeApiKey:
+		authToken = cred.ApiKey
+	case TypeToken:
+		p := jwt.NewParser()
+		tk, err := p.Parse(cred.GetAuthToken(), nil)
+		if err != nil && !errors.Is(err, jwt.ErrTokenUnverifiable) {
+			return fmt.Errorf("failed to parse oauth token: %w", err)
+		}
+
+		_, ok := tk.Claims.(jwt.MapClaims)["urn:columnar:trial_start"]
+		if !ok {
+			return ErrNoTrialLicense
+		}
+		authToken = cred.GetAuthToken()
+	default:
+		return fmt.Errorf("unsupported credential type: %s", cred.Type)
+	}
+
+	req, err := http.NewRequest(http.MethodGet, licenseURI, nil)
+	if err != nil {
+		return err
+	}
+
+	req.Header.Add("authorization", "Bearer "+authToken)
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		switch resp.StatusCode {
+		case http.StatusBadRequest:
+			return ErrNoTrialLicense
+		case http.StatusForbidden:
+			return ErrTrialExpired
+		default:
+			return fmt.Errorf("failed to fetch license: %s", resp.Status)
+		}
+	}
+
+	licenseFile, err := os.OpenFile(licensePath, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0o600)
+	if err != nil {
+		return err
+	}
+	defer licenseFile.Close()
+	if _, err = licenseFile.ReadFrom(resp.Body); err != nil {
+		licenseFile.Close()
+		os.Remove(licensePath)
+	}
+	return err
+}

cmd/dbc/auth.go

@@ -17,6 +17,7 @@ package main
 import (
 	"bufio"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/url"
@@ -83,6 +84,10 @@ func (l LoginCmd) GetModel() tea.Model {
 	)
 }
 
+type authSuccessMsg struct {
+	cred auth.Credential
+}
+
 type loginModel struct {
 	baseModel
 
@@ -190,13 +195,32 @@ func (m loginModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				}
 			})
 	case auth.Credential:
-		return m, tea.Sequence(func() tea.Msg {
+		return m, func() tea.Msg {
 			if err := auth.AddCredential(msg, true); err != nil {
 				return err
 			}
-			return nil
-		}, tea.Println("Authentication successful!"),
-			tea.Quit)
+			return authSuccessMsg{cred: msg}
+		}
+	case authSuccessMsg:
+		return m, tea.Sequence(tea.Println("Authentication successful!"),
+			func() tea.Msg {
+				if auth.IsColumnarPrivateRegistry((*url.URL)(&msg.cred.RegistryURL)) {
+					if err := auth.FetchColumnarLicense(&msg.cred); err != nil {
+						return err
+					}
+				}
+				return tea.Quit()
+			})
+	case error:
+		switch {
+		case errors.Is(msg, auth.ErrTrialExpired) ||
+			errors.Is(msg, auth.ErrNoTrialLicense):
+			// ignore these errors during auth login
+			// the user can still login but won't be able to download trial licenses
+			return m, tea.Quit
+		default:
+			// for other errors, let the baseModel update handle it.
+		}
 	}
 
 	base, cmd := m.baseModel.Update(msg)

cmd/dbc/main.go

@@ -15,6 +15,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"slices"
@@ -23,6 +24,7 @@ import (
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 	"github.com/columnar-tech/dbc"
+	"github.com/columnar-tech/dbc/auth"
 	"github.com/columnar-tech/dbc/cmd/dbc/completions"
 	"github.com/columnar-tech/dbc/config"
 	"github.com/mattn/go-isatty"
@@ -119,7 +121,16 @@ func (m baseModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 	case error:
 		m.status = 1
-		return m, tea.Sequence(tea.Println("Error: ", msg.Error()), tea.Quit)
+		var cmd tea.Cmd
+		switch {
+		case errors.Is(msg, auth.ErrTrialExpired):
+			cmd = tea.Println(errStyle.Render("Could not download license, trial has expired"))
+		case errors.Is(msg, auth.ErrNoTrialLicense):
+			cmd = tea.Println(errStyle.Render("Could not download license, trial not started"))
+		default:
+			cmd = tea.Println("Error: ", msg.Error())
+		}
+		return m, tea.Sequence(cmd, tea.Quit)
 	}
 	return m, nil
 }

drivers.go

@@ -164,6 +164,13 @@ func makereq(u string) (resp *http.Response, err error) {
 	}
 
 	if cred != nil {
+		if auth.IsColumnarPrivateRegistry(uri) {
+			// if we're accessing the private registry then attempt to
+			// fetch the trial license. This will be a no-op if they have
+			// a license saved already, and if they haven't started their
+			// trial or it is expired, then this will silently fail.
+			_ = auth.FetchColumnarLicense(cred)
+		}
 		req.Header.Set("Authorization", "Bearer "+cred.GetAuthToken())
 	}
 

go.mod

@@ -20,51 +20,54 @@ toolchain go1.24.6
 
 require (
 	github.com/Masterminds/semver/v3 v3.4.0
-	github.com/ProtonMail/gopenpgp/v3 v3.2.1
+	github.com/ProtonMail/gopenpgp/v3 v3.3.0
 	github.com/alexflint/go-arg v1.6.1
 	github.com/charmbracelet/bubbles v0.21.0
-	github.com/charmbracelet/bubbletea v1.3.7
+	github.com/charmbracelet/bubbletea v1.3.10
 	github.com/charmbracelet/lipgloss v1.1.0
 	github.com/cli/browser v1.3.0
-	github.com/cli/oauth v1.2.0
+	github.com/cli/oauth v1.2.1
 	github.com/go-faster/yaml v0.4.6
+	github.com/golang-jwt/jwt/v5 v5.3.0
 	github.com/google/uuid v1.6.0
 	github.com/mattn/go-isatty v0.0.20
 	github.com/pelletier/go-toml/v2 v2.2.4
-	github.com/stretchr/testify v1.10.0
-	github.com/zeroshade/machine-id v0.0.0-20250917170903-4283e98485ba
-	golang.org/x/sys v0.38.0
+	github.com/stretchr/testify v1.11.1
+	github.com/zeroshade/machine-id v0.0.0-20251223181436-930511047eef
+	golang.org/x/sys v0.39.0
 )
 
 require (
-	github.com/ProtonMail/go-crypto v1.2.0 // indirect
+	github.com/ProtonMail/go-crypto v1.3.0 // indirect
 	github.com/alexflint/go-scalar v1.2.0 // indirect
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
-	github.com/charmbracelet/colorprofile v0.3.1 // indirect
+	github.com/charmbracelet/colorprofile v0.4.1 // indirect
 	github.com/charmbracelet/harmonica v0.2.0 // indirect
-	github.com/charmbracelet/x/ansi v0.10.1 // indirect
-	github.com/charmbracelet/x/cellbuf v0.0.13 // indirect
-	github.com/charmbracelet/x/term v0.2.1 // indirect
-	github.com/cloudflare/circl v1.6.1 // indirect
-	github.com/columnar-tech/machine-id v0.0.0-20250917165521-f900c2b8afc9 // indirect
+	github.com/charmbracelet/x/ansi v0.11.3 // indirect
+	github.com/charmbracelet/x/cellbuf v0.0.14 // indirect
+	github.com/charmbracelet/x/term v0.2.2 // indirect
+	github.com/clipperhouse/displaywidth v0.6.2 // indirect
+	github.com/clipperhouse/stringish v0.1.1 // indirect
+	github.com/clipperhouse/uax29/v2 v2.3.0 // indirect
+	github.com/cloudflare/circl v1.6.2 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
 	github.com/go-faster/errors v0.7.1 // indirect
-	github.com/go-faster/jx v1.1.0 // indirect
-	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
+	github.com/go-faster/jx v1.2.0 // indirect
+	github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
 	github.com/mattn/go-localereader v0.0.1 // indirect
-	github.com/mattn/go-runewidth v0.0.16 // indirect
+	github.com/mattn/go-runewidth v0.0.19 // indirect
 	github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
 	github.com/muesli/cancelreader v0.2.2 // indirect
 	github.com/muesli/termenv v0.16.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rivo/uniseg v0.4.7 // indirect
 	github.com/sahilm/fuzzy v0.1.1 // indirect
-	github.com/segmentio/asm v1.2.0 // indirect
+	github.com/segmentio/asm v1.2.1 // indirect
 	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/crypto v0.45.0 // indirect
-	golang.org/x/text v0.31.0 // indirect
+	golang.org/x/crypto v0.46.0 // indirect
+	golang.org/x/text v0.32.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )