wrds-download

TUI/CLI tool for browsing and downloading WRDS data
Log | Files | Refs | README

commit 812b97abc3c38a81cc84b92ec6378b42a86c9e8e
parent a2b36d1f63d137ed45e6baaff0023c432a8aad56
Author: Erik Loualiche <[email protected]>
Date:   Fri, 20 Feb 2026 14:13:03 -0600

Fix TUI startup hang, Duo 2FA flooding, and Parquet compatibility

- Always start TUI in login mode so the user controls when the
  connection (and Duo 2FA push) fires — no more hanging on startup
- Limit pgxpool to MaxConns=1 to prevent multiple simultaneous
  auth attempts that trigger Duo security lockouts
- Add "Login with saved credentials" button to login form when
  credentials are available from config/env
- Force PLAIN encoding for Parquet string columns instead of
  DELTA_LENGTH_BYTE_ARRAY, which Julia's Parquet2.jl doesn't support

Co-Authored-By: Claude Opus 4.6 <[email protected]>

Diffstat:
Mcmd/tui.go | 24+++---------------------
Minternal/db/client.go | 10+++++++++-
Minternal/export/export.go | 1+
Minternal/tui/loginform.go | 192+++++++++++++++++++++++++++++++++++++++++++++++++++++--------------------------
4 files changed, 143 insertions(+), 84 deletions(-)

diff --git a/cmd/tui.go b/cmd/tui.go @@ -1,13 +1,11 @@ package cmd import ( - "context" "fmt" "os" tea "github.com/charmbracelet/bubbletea" "github.com/eloualiche/wrds-download/internal/config" - "github.com/eloualiche/wrds-download/internal/db" "github.com/eloualiche/wrds-download/internal/tui" "github.com/spf13/cobra" ) @@ -25,25 +23,9 @@ func init() { func runTUI(cmd *cobra.Command, args []string) error { config.ApplyCredentials() - ctx := context.Background() - client, err := db.New(ctx) - if err != nil { - // Launch TUI in login mode instead of crashing - m := tui.NewAppNoClient() - p := tea.NewProgram(m, tea.WithAltScreen()) - final, err := p.Run() - if err != nil { - fmt.Fprintln(os.Stderr, err) - os.Exit(1) - } - if a, ok := final.(*tui.App); ok && a.Err() != "" { - fmt.Fprintln(os.Stderr, a.Err()) - } - return nil - } - defer client.Close() - - m := tui.NewApp(client) + // Always start in login mode so the user controls when the + // connection (and any 2FA prompt) happens. + m := tui.NewAppNoClient() p := tea.NewProgram(m, tea.WithAltScreen()) final, err := p.Run() if err != nil { diff --git a/internal/db/client.go b/internal/db/client.go @@ -48,12 +48,20 @@ func getenv(key, fallback string) string { } // New creates and pings a pgx pool using DSNFromEnv. +// The pool is limited to a single connection to avoid triggering +// multiple authentication prompts (e.g. Duo 2FA on WRDS). func New(ctx context.Context) (*Client, error) { dsn, err := DSNFromEnv() if err != nil { return nil, err } - pool, err := pgxpool.New(ctx, dsn) + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + return nil, fmt.Errorf("parse dsn: %w", err) + } + cfg.MaxConns = 1 + cfg.MinConns = 0 + pool, err := pgxpool.NewWithConfig(ctx, cfg) if err != nil { return nil, fmt.Errorf("pgxpool.New: %w", err) } diff --git a/internal/export/export.go b/internal/export/export.go @@ -119,6 +119,7 @@ func writeParquet(rows pgx.Rows, outPath string) error { writer := parquet.NewGenericWriter[map[string]any](f, schema, parquet.Compression(&zstd.Codec{}), + parquet.DefaultEncodingFor(parquet.ByteArray, &parquet.Plain), ) buf := make([]map[string]any, 0, rowGroupSize) diff --git a/internal/tui/loginform.go b/internal/tui/loginform.go @@ -1,6 +1,7 @@ package tui import ( + "os" "strings" "github.com/charmbracelet/bubbles/textinput" @@ -11,20 +12,24 @@ import ( type loginField int const ( - loginFieldUser loginField = iota + loginFieldSaved loginField = iota // "Login with saved credentials" button + loginFieldUser loginFieldPassword loginFieldDatabase loginFieldSave loginFieldCount ) -const loginTextInputs = 3 // number of text input fields (before the save toggle) +const loginTextInputs = 3 // user, password, database // LoginForm is the login dialog overlay shown when credentials are missing. type LoginForm struct { - inputs [loginTextInputs]textinput.Model - save bool - focused loginField + inputs [loginTextInputs]textinput.Model + save bool + focused loginField + savedUser string // non-empty when saved credentials are available + savedPw string + savedDB string } // LoginSubmitMsg is sent when the user confirms the login form. @@ -41,26 +46,95 @@ type LoginCancelMsg struct{} func newLoginForm() LoginForm { f := LoginForm{} - f.inputs[loginFieldUser] = textinput.New() - f.inputs[loginFieldUser].Placeholder = "WRDS username" - f.inputs[loginFieldUser].CharLimit = 128 + f.inputs[loginFieldUser-1] = textinput.New() + f.inputs[loginFieldUser-1].Placeholder = "WRDS username" + f.inputs[loginFieldUser-1].CharLimit = 128 - f.inputs[loginFieldPassword] = textinput.New() - f.inputs[loginFieldPassword].Placeholder = "WRDS password" - f.inputs[loginFieldPassword].CharLimit = 128 - f.inputs[loginFieldPassword].EchoMode = textinput.EchoPassword - f.inputs[loginFieldPassword].EchoCharacter = '*' + f.inputs[loginFieldPassword-1] = textinput.New() + f.inputs[loginFieldPassword-1].Placeholder = "WRDS password" + f.inputs[loginFieldPassword-1].CharLimit = 128 + f.inputs[loginFieldPassword-1].EchoMode = textinput.EchoPassword + f.inputs[loginFieldPassword-1].EchoCharacter = '*' - f.inputs[loginFieldDatabase] = textinput.New() - f.inputs[loginFieldDatabase].Placeholder = "wrds" - f.inputs[loginFieldDatabase].CharLimit = 128 - f.inputs[loginFieldDatabase].SetValue("wrds") + f.inputs[loginFieldDatabase-1] = textinput.New() + f.inputs[loginFieldDatabase-1].Placeholder = "wrds" + f.inputs[loginFieldDatabase-1].CharLimit = 128 + f.inputs[loginFieldDatabase-1].SetValue("wrds") + + // Check for saved credentials in env (set by config.ApplyCredentials). + f.savedUser = os.Getenv("PGUSER") + f.savedPw = os.Getenv("PGPASSWORD") + f.savedDB = os.Getenv("PGDATABASE") + if f.savedDB == "" { + f.savedDB = "wrds" + } f.save = true - f.inputs[loginFieldUser].Focus() + + if f.hasSaved() { + f.focused = loginFieldSaved + } else { + f.focused = loginFieldUser + f.inputs[loginFieldUser-1].Focus() + } return f } +func (f *LoginForm) hasSaved() bool { + return f.savedUser != "" && f.savedPw != "" +} + +// inputIndex maps a loginField to the inputs array index. +// Returns -1 for non-input fields (saved, save). +func inputIndex(field loginField) int { + switch field { + case loginFieldUser, loginFieldPassword, loginFieldDatabase: + return int(field) - 1 + } + return -1 +} + +func (f *LoginForm) blurCurrent() { + if idx := inputIndex(f.focused); idx >= 0 { + f.inputs[idx].Blur() + } +} + +func (f *LoginForm) focusCurrent() tea.Cmd { + if idx := inputIndex(f.focused); idx >= 0 { + f.inputs[idx].Focus() + return textinput.Blink + } + return nil +} + +func (f *LoginForm) advance(delta int) tea.Cmd { + f.blurCurrent() + start := loginFieldSaved + if !f.hasSaved() { + start = loginFieldUser + } + count := int(loginFieldCount) - int(start) + pos := (int(f.focused) - int(start) + delta%count + count) % count + f.focused = loginField(pos + int(start)) + return f.focusCurrent() +} + +func (f LoginForm) submit() tea.Cmd { + user := strings.TrimSpace(f.inputs[loginFieldUser-1].Value()) + pw := f.inputs[loginFieldPassword-1].Value() + if user == "" || pw == "" { + return nil + } + database := strings.TrimSpace(f.inputs[loginFieldDatabase-1].Value()) + if database == "" { + database = "wrds" + } + return func() tea.Msg { + return LoginSubmitMsg{User: user, Password: pw, Database: database, Save: f.save} + } +} + func (f LoginForm) Update(msg tea.Msg) (LoginForm, tea.Cmd) { switch msg := msg.(type) { case tea.KeyMsg: @@ -69,53 +143,31 @@ func (f LoginForm) Update(msg tea.Msg) (LoginForm, tea.Cmd) { return f, func() tea.Msg { return LoginCancelMsg{} } case "enter": - if f.focused == loginFieldSave { - // Submit - user := strings.TrimSpace(f.inputs[loginFieldUser].Value()) - pw := f.inputs[loginFieldPassword].Value() - if user == "" || pw == "" { - return f, nil - } - database := strings.TrimSpace(f.inputs[loginFieldDatabase].Value()) - if database == "" { - database = "wrds" - } + if f.focused == loginFieldSaved { + // Connect using saved credentials. return f, func() tea.Msg { - return LoginSubmitMsg{User: user, Password: pw, Database: database, Save: f.save} + return LoginSubmitMsg{ + User: f.savedUser, + Password: f.savedPw, + Database: f.savedDB, + Save: false, + } } } - // Advance to next field - if int(f.focused) < loginTextInputs { - f.inputs[f.focused].Blur() - } - f.focused++ - if int(f.focused) < loginTextInputs { - f.inputs[f.focused].Focus() - return f, textinput.Blink + if f.focused == loginFieldSave { + return f, f.submit() } - return f, nil + // Advance to next field. + cmd := f.advance(1) + return f, cmd case "tab", "down": - if int(f.focused) < loginTextInputs { - f.inputs[f.focused].Blur() - } - f.focused = loginField((int(f.focused) + 1) % int(loginFieldCount)) - if int(f.focused) < loginTextInputs { - f.inputs[f.focused].Focus() - return f, textinput.Blink - } - return f, nil + cmd := f.advance(1) + return f, cmd case "shift+tab", "up": - if int(f.focused) < loginTextInputs { - f.inputs[f.focused].Blur() - } - f.focused = loginField((int(f.focused) + int(loginFieldCount) - 1) % int(loginFieldCount)) - if int(f.focused) < loginTextInputs { - f.inputs[f.focused].Focus() - return f, textinput.Blink - } - return f, nil + cmd := f.advance(-1) + return f, cmd case " ": if f.focused == loginFieldSave { @@ -125,10 +177,10 @@ func (f LoginForm) Update(msg tea.Msg) (LoginForm, tea.Cmd) { } } - // Forward to focused text input - if int(f.focused) < loginTextInputs { + // Forward to focused text input. + if idx := inputIndex(f.focused); idx >= 0 { var cmd tea.Cmd - f.inputs[f.focused], cmd = f.inputs[f.focused].Update(msg) + f.inputs[idx], cmd = f.inputs[idx].Update(msg) return f, cmd } return f, nil @@ -140,17 +192,33 @@ func (f LoginForm) View(width int, errMsg string) string { title := stylePanelHeader.Render("WRDS Login") sb.WriteString(title + "\n\n") + // "Login with saved credentials" button. + if f.hasSaved() { + btnLabel := "Login as " + f.savedUser + btnStyle := lipgloss.NewStyle().Foreground(colorMuted) + if f.focused == loginFieldSaved { + btnStyle = lipgloss.NewStyle(). + Foreground(lipgloss.Color("#FFFFFF")). + Background(colorFocus). + Padding(0, 1) + btnLabel += " [enter]" + } + sb.WriteString(btnStyle.Render(btnLabel) + "\n\n") + sb.WriteString(lipgloss.NewStyle().Foreground(colorMuted).Render("── or enter credentials manually ──") + "\n\n") + } + labels := []string{"Username", "Password", "Database"} + fields := []loginField{loginFieldUser, loginFieldPassword, loginFieldDatabase} for i, label := range labels { style := lipgloss.NewStyle().Foreground(colorMuted) - if loginField(i) == f.focused { + if fields[i] == f.focused { style = lipgloss.NewStyle().Foreground(colorFocus) } sb.WriteString(style.Render(label+" ") + "\n") sb.WriteString(f.inputs[i].View() + "\n\n") } - // Save toggle + // Save toggle. check := "[ ]" if f.save { check = "[x]"