diff --git a/.gitignore b/.gitignore index b51d491c..d3b19d12 100644 --- a/.gitignore +++ b/.gitignore @@ -54,4 +54,6 @@ docs/uv.lock *.code-workspace *.vscode/ *.claude/ + +# Python __pycache__/ diff --git a/cmd/dotenv/if_needed_test.go b/cmd/dotenv/if_needed_test.go index d4c004e7..d269ed6e 100644 --- a/cmd/dotenv/if_needed_test.go +++ b/cmd/dotenv/if_needed_test.go @@ -22,6 +22,15 @@ import ( "github.com/stretchr/testify/require" ) +func TestMain(m *testing.M) { + // Clear DataRobot environment variables that might leak from the host environment + // into tests, ensuring tests have a clean isolated environment. + os.Unsetenv("DATAROBOT_ENDPOINT") + os.Unsetenv("DATAROBOT_API_TOKEN") + + os.Exit(m.Run()) +} + func TestShouldSkipSetup(t *testing.T) { t.Run("should skip when .env exists and validation passes", func(t *testing.T) { tmpDir := t.TempDir() diff --git a/cmd/start/model.go b/cmd/start/model.go index a147bb83..2409dffe 100644 --- a/cmd/start/model.go +++ b/cmd/start/model.go @@ -56,6 +56,7 @@ type Model struct { waitingToExecute bool // Whether to wait for user input before proceeding needTemplateSetup bool // Whether we need to run template setup after quitting repoRoot string + pulumiLoginModel *pulumiLoginModel // Sub-model for Pulumi login flow } type stepCompleteMsg struct { @@ -67,6 +68,7 @@ type stepCompleteMsg struct { selfUpdate bool // Whether to ask for self update executeScript bool // Whether to execute the script immediately needTemplateSetup bool // Whether we need to run template setup + needPulumiLogin bool // Whether we need to enter Pulumi login flow } type startScriptCompleteMsg struct{ err error } @@ -97,6 +99,7 @@ func NewStartModel(opts Options) Model { // TODO Implement validateEnvironment // {description: "Validating environment...", fn: validateEnvironment}, {description: "Checking repository setup...", fn: checkRepository}, + {description: "Checking Pulumi state backend...", fn: checkPulumiState}, {description: "Finding and executing start command...", fn: findAndExecuteStart}, }, opts: opts, @@ -186,6 +189,11 @@ func (m Model) execSelfUpdate() tea.Cmd { } func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + // If Pulumi login sub-model is active, delegate to it + if m.pulumiLoginModel != nil { + return m.handlePulumiLoginUpdate(msg) + } + switch msg := msg.(type) { case tea.KeyMsg: return m.handleKey(msg) @@ -220,6 +228,29 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, nil } +func (m Model) handlePulumiLoginUpdate(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case pulumiLoginCompleteMsg: + // Pulumi login completed, clear sub-model and continue + m.pulumiLoginModel = nil + return m.executeNextStep() + + case pulumiLoginErrorMsg: + // Pulumi login failed + m.err = msg.err + return m, tea.Quit + + default: + // Delegate to sub-model + subModel, cmd := m.pulumiLoginModel.Update(msg) + if plm, ok := subModel.(pulumiLoginModel); ok { + m.pulumiLoginModel = &plm + } + + return m, cmd + } +} + func (m Model) handleKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { // If there's an error, any key press quits if m.err != nil { @@ -289,29 +320,17 @@ func (m Model) handleStepComplete(msg stepCompleteMsg) (tea.Model, tea.Cmd) { "execute_script", msg.executeScript, "quickstart_script_path", msg.quickstartScriptPath, "need_template_setup", msg.needTemplateSetup, + "need_pulumi_login", msg.needPulumiLogin, ) - // Store any message from the completed step - if msg.message != "" { - m.stepCompleteMessage = msg.message - } + m.updateFromStepComplete(msg) - if msg.hideMenu { - m.hideMenu = msg.hideMenu - } + // If we need to enter Pulumi login flow + if msg.needPulumiLogin { + plm := newPulumiLoginModel() + m.pulumiLoginModel = &plm - if msg.selfUpdate { - m.selfUpdate = msg.selfUpdate - } - - // Store quickstart script path if provided - if msg.quickstartScriptPath != "" { - m.quickstartScriptPath = msg.quickstartScriptPath - } - - // Store whether we need template setup - if msg.needTemplateSetup { - m.needTemplateSetup = true + return m, plm.Init() } // If this step requires executing a script, do it now @@ -336,9 +355,57 @@ func (m Model) handleStepComplete(msg stepCompleteMsg) (tea.Model, tea.Cmd) { return m.executeNextStep() } +func (m *Model) updateFromStepComplete(msg stepCompleteMsg) { + // Store any message from the completed step + if msg.message != "" { + m.stepCompleteMessage = msg.message + } + + if msg.hideMenu { + m.hideMenu = msg.hideMenu + } + + if msg.selfUpdate { + m.selfUpdate = msg.selfUpdate + } + + // Store quickstart script path if provided + if msg.quickstartScriptPath != "" { + m.quickstartScriptPath = msg.quickstartScriptPath + } + + // Store whether we need template setup + if msg.needTemplateSetup { + m.needTemplateSetup = true + } +} + func (m Model) View() string { //nolint: cyclop var sb strings.Builder + // If Pulumi login sub-model is active, show it + if m.pulumiLoginModel != nil { + sb.WriteString("\n") + sb.WriteString(tui.WelcomeStyle.Render("🚀 DataRobot AI Application Quickstart")) + sb.WriteString("\n\n") + + // Show progress through steps + for i, step := range m.steps { + if i < m.current { + sb.WriteString(fmt.Sprintf(" %s %s\n", checkMark, tui.DimStyle.Render(step.description))) + } else if i == m.current { + sb.WriteString(fmt.Sprintf(" %s %s\n", arrow, step.description)) + } else { + sb.WriteString(fmt.Sprintf(" %s\n", tui.DimStyle.Render(step.description))) + } + } + + sb.WriteString("\n") + sb.WriteString(m.pulumiLoginModel.View()) + + return sb.String() + } + if !m.hideMenu { sb.WriteString("\n") sb.WriteString(tui.WelcomeStyle.Render("🚀 DataRobot AI Application Quickstart")) diff --git a/cmd/start/pulumiModel.go b/cmd/start/pulumiModel.go new file mode 100644 index 00000000..d8e831c8 --- /dev/null +++ b/cmd/start/pulumiModel.go @@ -0,0 +1,359 @@ +// Copyright 2025 DataRobot, Inc. and its affiliates. +// +// 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 start + +import ( + "crypto/rand" + "encoding/base64" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + + "github.com/charmbracelet/bubbles/textinput" + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + "github.com/datarobot/cli/tui" + "github.com/spf13/viper" +) + +const ( + generatedPassphraseLength = 32 + pulumiConfigPassphraseKey = "pulumi_config_passphrase" + pulumiDocsURL = "https://www.pulumi.com/docs/iac/concepts/state-and-backends/" +) + +type pulumiLoginScreen int + +const ( + pulumiLoginScreenBackendSelection pulumiLoginScreen = iota + pulumiLoginScreenDIYURL + pulumiLoginScreenPassphrasePrompt + pulumiLoginScreenLoggingIn +) + +// pulumiLoginModel handles the Pulumi login flow +type pulumiLoginModel struct { + currentScreen pulumiLoginScreen + selectedOption int + options []string + diyInput textinput.Model + diyURL string + wantsPassphrase bool + generatedPassphrase string + err error + loginOutput string +} + +type ( + pulumiLoginCompleteMsg struct{} + pulumiLoginErrorMsg struct{ err error } + pulumiLoginSuccessMsg struct{ output string } +) + +func newPulumiLoginModel() pulumiLoginModel { + ti := textinput.New() + ti.Placeholder = "s3://my-pulumi-bucket or azblob://..." + ti.Focus() + ti.Width = 60 + + return pulumiLoginModel{ + currentScreen: pulumiLoginScreenBackendSelection, + selectedOption: 0, + options: []string{"Login locally", "Login to Pulumi Cloud", "DIY backend (S3, Azure Blob, etc.)"}, + diyInput: ti, + } +} + +func (m pulumiLoginModel) Init() tea.Cmd { + return textinput.Blink +} + +func (m pulumiLoginModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.KeyMsg: + return m.handleKeyPress(msg) + + case pulumiLoginSuccessMsg: + m.loginOutput = msg.output + return m, func() tea.Msg { return pulumiLoginCompleteMsg{} } + + case pulumiLoginErrorMsg: + m.err = msg.err + return m, nil + } + + // Handle text input updates for DIY URL screen + if m.currentScreen == pulumiLoginScreenDIYURL { + var cmd tea.Cmd + + m.diyInput, cmd = m.diyInput.Update(msg) + + return m, cmd + } + + return m, nil +} + +func (m pulumiLoginModel) handleKeyPress(msg tea.KeyMsg) (tea.Model, tea.Cmd) { + switch m.currentScreen { + case pulumiLoginScreenBackendSelection: + return m.handleBackendSelectionKey(msg) + + case pulumiLoginScreenDIYURL: + return m.handleDIYURLKey(msg) + + case pulumiLoginScreenPassphrasePrompt: + return m.handlePassphrasePromptKey(msg) + + case pulumiLoginScreenLoggingIn: + // No key handling during login + return m, nil + } + + return m, nil +} + +func (m pulumiLoginModel) handleBackendSelectionKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { + switch msg.String() { + case "up", "k": + if m.selectedOption > 0 { + m.selectedOption-- + } + case "down", "j": + if m.selectedOption < len(m.options)-1 { + m.selectedOption++ + } + case "enter": + // User selected an option + switch m.selectedOption { + case 0: // Local + m.currentScreen = pulumiLoginScreenPassphrasePrompt + case 1: // Cloud + return m, m.performLogin("cloud", "") + case 2: // DIY + m.currentScreen = pulumiLoginScreenDIYURL + } + } + + return m, nil +} + +func (m pulumiLoginModel) handleDIYURLKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { + switch msg.String() { + case "enter": + m.diyURL = strings.TrimSpace(m.diyInput.Value()) + if m.diyURL == "" { + return m, nil + } + + m.currentScreen = pulumiLoginScreenPassphrasePrompt + case "esc": + m.currentScreen = pulumiLoginScreenBackendSelection + default: + var cmd tea.Cmd + + m.diyInput, cmd = m.diyInput.Update(msg) + + return m, cmd + } + + return m, nil +} + +func (m pulumiLoginModel) handlePassphrasePromptKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { + switch msg.String() { + case "y", "Y": + m.wantsPassphrase = true + // Generate passphrase + passphrase, err := generateRandomPassphrase(generatedPassphraseLength) + if err != nil { + m.err = fmt.Errorf("failed to generate passphrase: %w", err) + return m, nil + } + + m.generatedPassphrase = passphrase + + // Perform login + switch m.selectedOption { + case 0: // Local + return m, m.performLogin("local", "") + case 2: // DIY + return m, m.performLogin("diy", m.diyURL) + } + case "n", "N": + m.wantsPassphrase = false + // Perform login without passphrase + switch m.selectedOption { + case 0: // Local + return m, m.performLogin("local", "") + case 2: // DIY + return m, m.performLogin("diy", m.diyURL) + } + case "esc": + // Go back + if m.selectedOption == 2 { // DIY + m.currentScreen = pulumiLoginScreenDIYURL + } else { + m.currentScreen = pulumiLoginScreenBackendSelection + } + } + + return m, nil +} + +func (m pulumiLoginModel) savePassphraseToConfig() error { + homeDir, err := os.UserHomeDir() + if err != nil { + return fmt.Errorf("failed to get home directory: %w", err) + } + + configDir := filepath.Join(homeDir, ".config", "datarobot") + if err := os.MkdirAll(configDir, os.ModePerm); err != nil { + return fmt.Errorf("failed to create config directory: %w", err) + } + + viper.Set(pulumiConfigPassphraseKey, m.generatedPassphrase) + + if err := viper.WriteConfig(); err != nil { + return fmt.Errorf("failed to write config: %w", err) + } + + return nil +} + +func (m pulumiLoginModel) performLogin(loginType, url string) tea.Cmd { + return func() tea.Msg { + var cmd *exec.Cmd + + // Set passphrase in config if requested + if m.wantsPassphrase && m.generatedPassphrase != "" { + if err := m.savePassphraseToConfig(); err != nil { + return pulumiLoginErrorMsg{err} + } + } + + // Determine which pulumi login command to run + switch loginType { + case "local": + cmd = exec.Command("pulumi", "login", "--local") + case "cloud": + cmd = exec.Command("pulumi", "login") + case "diy": + cmd = exec.Command("pulumi", "login", url) + default: + return pulumiLoginErrorMsg{errors.New("unknown login type")} + } + + // Run the command + output, err := cmd.CombinedOutput() + if err != nil { + return pulumiLoginErrorMsg{fmt.Errorf("pulumi login failed: %w\n%s", err, string(output))} + } + + return pulumiLoginSuccessMsg{output: string(output)} + } +} + +func (m pulumiLoginModel) View() string { + var sb strings.Builder + + if m.err != nil { + sb.WriteString(tui.ErrorStyle.Render("Error: ") + m.err.Error() + "\n") + return sb.String() + } + + switch m.currentScreen { + case pulumiLoginScreenBackendSelection: + sb.WriteString(tui.SubTitleStyle.Render("Pulumi State Backend Selection")) + sb.WriteString("\n\n") + sb.WriteString("Select where Pulumi should store your infrastructure state:\n\n") + + for i, option := range m.options { + cursor := " " + if i == m.selectedOption { + cursor = arrow.String() + " " + } + + sb.WriteString(fmt.Sprintf("%s%s\n", cursor, option)) + } + + sb.WriteString("\n") + sb.WriteString(tui.DimStyle.Render("↑/↓ to navigate • enter to select")) + + case pulumiLoginScreenDIYURL: + sb.WriteString(tui.SubTitleStyle.Render("DIY Backend Configuration")) + sb.WriteString("\n\n") + sb.WriteString(fmt.Sprintf("For more information about backends, see:\n%s\n\n", + lipgloss.NewStyle().Foreground(tui.DrPurple).Render(pulumiDocsURL))) + sb.WriteString("Enter your backend URL:\n") + sb.WriteString("Examples: s3://my-pulumi-bucket, azblob://..., gs://...\n\n") + sb.WriteString(m.diyInput.View()) + sb.WriteString("\n\n") + sb.WriteString(tui.DimStyle.Render("enter to continue • esc to go back")) + + case pulumiLoginScreenPassphrasePrompt: + sb.WriteString(tui.SubTitleStyle.Render("Pulumi Configuration Passphrase")) + sb.WriteString("\n\n") + sb.WriteString("Would you like to set a default PULUMI_CONFIG_PASSPHRASE?\n") + sb.WriteString("This will be used to encrypt secrets and stack variables.\n\n") + sb.WriteString("We can auto-generate a strong passphrase and save it to your\n") + sb.WriteString("DataRobot CLI config file (~/.config/datarobot/drconfig.yaml)\n\n") + sb.WriteString(tui.DimStyle.Render("y to generate passphrase • n to skip • esc to go back")) + + case pulumiLoginScreenLoggingIn: + sb.WriteString("Logging in to Pulumi...\n") + } + + return sb.String() +} + +func generateRandomPassphrase(length int) (string, error) { + bytes := make([]byte, length) + + _, err := rand.Read(bytes) + if err != nil { + return "", fmt.Errorf("failed to generate random bytes: %w", err) + } + + return base64.URLEncoding.EncodeToString(bytes)[:length], nil +} + +// checkPulumiState checks if user is logged into Pulumi and guides them through login if needed +func checkPulumiState(m *Model) tea.Msg { + // Check if pulumi is installed + _, err := exec.LookPath("pulumi") + if err != nil { + // Pulumi not installed, skip this step + return stepCompleteMsg{} + } + + // Check if user is logged in using 'pulumi whoami' + cmd := exec.Command("pulumi", "whoami") + + err = cmd.Run() + if err == nil { + // User is already logged in + return stepCompleteMsg{} + } + + // User is not logged in, need to guide them through login + // Signal to the main model to enter Pulumi login sub-model + return stepCompleteMsg{ + needPulumiLogin: true, + } +} diff --git a/cmd/start/pulumiModel_test.go b/cmd/start/pulumiModel_test.go new file mode 100644 index 00000000..51bfdaec --- /dev/null +++ b/cmd/start/pulumiModel_test.go @@ -0,0 +1,56 @@ +// Copyright 2025 DataRobot, Inc. and its affiliates. +// +// 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 start + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCheckPulumiState_PulumiNotInstalled(t *testing.T) { + // This test assumes pulumi is not in PATH with a name like "pulumi-nonexistent" + // The checkPulumiState function should skip the check gracefully + m := &Model{} + msg := checkPulumiState(m) + + // When pulumi is not installed, should return stepCompleteMsg with defaults + stepMsg, ok := msg.(stepCompleteMsg) + assert.True(t, ok, "Expected stepCompleteMsg") + assert.False(t, stepMsg.needPulumiLogin, "Should not need Pulumi login when not installed") +} + +func TestPulumiLoginModel_InitialState(t *testing.T) { + model := newPulumiLoginModel() + + assert.Equal(t, pulumiLoginScreenBackendSelection, model.currentScreen) + assert.Equal(t, 0, model.selectedOption) + assert.Len(t, model.options, 3) + assert.Equal(t, "Login locally", model.options[0]) + assert.Equal(t, "Login to Pulumi Cloud", model.options[1]) + assert.Contains(t, model.options[2], "DIY") +} + +func TestGenerateRandomPassphrase(t *testing.T) { + passphrase, err := generateRandomPassphrase(32) + require.NoError(t, err) + assert.Len(t, passphrase, 32) + + // Generate another one to ensure they're different + passphrase2, err := generateRandomPassphrase(32) + require.NoError(t, err) + assert.NotEqual(t, passphrase, passphrase2, "Generated passphrases should be unique") +} diff --git a/internal/envbuilder/builder.go b/internal/envbuilder/builder.go index 1473b671..ccfb17c0 100644 --- a/internal/envbuilder/builder.go +++ b/internal/envbuilder/builder.go @@ -148,10 +148,6 @@ func GatherUserPrompts(rootDir string, variables Variables) ([]UserPrompt, error return nil, fmt.Errorf("Failed to discover task yaml files: %w", err) } - if len(yamlFiles) == 0 { - return nil, nil - } - allPrompts := make([]UserPrompt, 0) allPrompts = append(allPrompts, corePrompts...) diff --git a/internal/envbuilder/validator.go b/internal/envbuilder/validator.go index 350ea135..7e1d2d46 100644 --- a/internal/envbuilder/validator.go +++ b/internal/envbuilder/validator.go @@ -18,6 +18,8 @@ import ( "os" "slices" "strings" + + "github.com/spf13/viper" ) // ValidationResult represents the validation status of a single variable. @@ -109,11 +111,30 @@ func ValidateEnvironment(repoRoot string, variables Variables) EnvironmentValida // promptsWithValues updates slice of prompts with values from .env file contents // and environment variables (environment variables take precedence). func promptsWithValues(prompts []UserPrompt, variables Variables) []UserPrompt { + // Special handling for PULUMI_CONFIG_PASSPHRASE from viper config + // This happens regardless of whether variables exist + for p := range prompts { + if prompts[p].Env == "PULUMI_CONFIG_PASSPHRASE" { + // Check if already set in environment + if _, ok := os.LookupEnv(prompts[p].Env); !ok { + // Not in environment, check viper config + if configValue := viper.GetString("pulumi_config_passphrase"); configValue != "" { + prompts[p].Value = configValue + } + } + } + } + if len(variables) == 0 { return prompts } for p, prompt := range prompts { + // Skip if already processed (like PULUMI_CONFIG_PASSPHRASE above) + if prompt.Value != "" { + continue + } + // Capture existing env var values if existingEnvValue, ok := os.LookupEnv(prompt.Env); ok { prompt.Value = existingEnvValue diff --git a/internal/envbuilder/validator_test.go b/internal/envbuilder/validator_test.go index fc0ffc46..1f433f22 100644 --- a/internal/envbuilder/validator_test.go +++ b/internal/envbuilder/validator_test.go @@ -183,6 +183,25 @@ func TestPromptsWithValues(t *testing.T) { t.Errorf("Expected TEST_OVERRIDE to be overridden to 'from-env', got '%s'", result[0].Value) } }) + t.Run("PULUMI_CONFIG_PASSPHRASE reads from viper config", func(t *testing.T) { + // This test verifies that PULUMI_CONFIG_PASSPHRASE gets its value from viper config + // when it's not in environment or .env file + // Note: In real usage, viper would be initialized with the config file + // For this test, we're just verifying the code path exists and falls back gracefully + variables := Variables{} + + prompts := []UserPrompt{ + {Env: "PULUMI_CONFIG_PASSPHRASE", Default: "default-pass"}, + } + + result := promptsWithValues(prompts, variables) + + // When variables is empty and viper config is not set, should remain empty + // This allows proper validation - the value will be filled from viper when it exists + if result[0].Value != "" { + t.Errorf("Expected PULUMI_CONFIG_PASSPHRASE to be empty (for validation), got '%s'", result[0].Value) + } + }) } func TestIsOptionSelected(t *testing.T) {