Skip to content

Commit

Permalink
Merge pull request #205 from bruin-data/patch/manage-uv-under-bruin
Browse files Browse the repository at this point in the history 
Patch/manage uv under bruin
  • Loading branch information
@albertobruin
albertobruin authored Nov 13, 2024
2 parents 677a93a + 4777fd5 commit f729c14
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 38 deletions.
2 changes: 1 addition & 1 deletion examples/simple-pipeline/assets/ingestr_duckdb_asset.yml
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
name: pulic.postgres_to_duckdb
name: public.postgres_to_duckdb
type: ingestr

parameters:
Expand Down
75 changes: 46 additions & 29 deletions pkg/python/uv.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,9 @@ import (
"github.com/bruin-data/bruin/pkg/executor"
"github.com/bruin-data/bruin/pkg/git"
"github.com/bruin-data/bruin/pkg/pipeline"
"github.com/bruin-data/bruin/pkg/user"
"github.com/pkg/errors"
"github.com/spf13/afero"
)

var AvailablePythonVersions = map[string]bool{
Expand All @@ -42,46 +44,57 @@ type UvChecker struct {
cmd CommandRunner
}

// EnsureUvInstalled checks if uv is installed and installs it if not present.
func (u *UvChecker) EnsureUvInstalled(ctx context.Context) error {
// EnsureUvInstalled checks if uv is installed and installs it if not present, then returns the full path of the binary.
func (u *UvChecker) EnsureUvInstalled(ctx context.Context) (string, error) {
u.mut.Lock()
defer u.mut.Unlock()

// Check if uv is already installed
_, err := exec.LookPath("uv")
m := user.NewConfigManager(afero.NewOsFs())
bruinHomeDirAbsPath, err := m.EnsureAndGetBruinHomeDir()
if err != nil {
err = u.installUvCommand(ctx)
return "", errors.Wrap(err, "failed to get bruin home directory")
}
var binaryName string
if runtime.GOOS == "windows" {
binaryName = "uv.exe"
} else {
binaryName = "uv"
}
uvBinaryPath := filepath.Join(bruinHomeDirAbsPath, binaryName)
if _, err := os.Stat(uvBinaryPath); errors.Is(err, os.ErrNotExist) {
err = u.installUvCommand(ctx, bruinHomeDirAbsPath)
if err != nil {
return err
return "", err
}
return nil
return uvBinaryPath, nil
}

cmd := exec.Command("uv", "version", "--output-format", "json")
cmd := exec.Command(uvBinaryPath, "version", "--output-format", "json")
output, err := cmd.CombinedOutput()
if err != nil {
return fmt.Errorf("failed to check uv version: %w\nOutput: %s\n\nPlease install uv v%s yourself: https://docs.astral.sh/uv/getting-started/installation/", err, output, UvVersion)
return "", fmt.Errorf("failed to check uv version: %w -- Output: %s", err, output)
}

var uvVersion struct {
Version string `json:"version"`
}
if err := json.Unmarshal(output, &uvVersion); err != nil {
return fmt.Errorf("failed to parse uv version: %w", err)
return "", fmt.Errorf("failed to parse uv version: %w", err)
}

if uvVersion.Version != UvVersion {
err = u.installUvCommand(ctx)
err = u.installUvCommand(ctx, bruinHomeDirAbsPath)
if err != nil {
return err
return "", err
}
return nil
return uvBinaryPath, nil
}

return nil
return uvBinaryPath, nil
}

func (u *UvChecker) installUvCommand(ctx context.Context) error {
func (u *UvChecker) installUvCommand(ctx context.Context, dest string) error {
var output io.Writer = os.Stdout
if ctx.Value(executor.KeyPrinter) != nil {
output = ctx.Value(executor.KeyPrinter).(io.Writer)
Expand All @@ -94,14 +107,14 @@ func (u *UvChecker) installUvCommand(ctx context.Context) error {

var commandInstance *exec.Cmd
if runtime.GOOS == "windows" {
commandInstance = exec.Command(Shell, ShellSubcommandFlag, fmt.Sprintf("winget install --accept-package-agreements --accept-source-agreements --silent --id=astral-sh.uv --version %s -e", UvVersion)) //nolint:gosec
commandInstance = exec.Command(Shell, ShellSubcommandFlag, fmt.Sprintf("winget install --accept-package-agreements --accept-source-agreements --silent --id=astral-sh.uv --version %s --location %s -e", UvVersion, dest)) //nolint:gosec
} else {
commandInstance = exec.Command(Shell, ShellSubcommandFlag, fmt.Sprintf(" set -o pipefail; curl -LsSf https://astral.sh/uv/%s/install.sh | sh", UvVersion)) //nolint:gosec
commandInstance = exec.Command(Shell, ShellSubcommandFlag, fmt.Sprintf(" set -o pipefail; curl -LsSf https://astral.sh/uv/%s/install.sh | UV_INSTALL_DIR=\"%s\" NO_MODIFY_PATH=1 sh", UvVersion, dest)) //nolint:gosec
}

err := u.cmd.RunAnyCommand(ctx, commandInstance)
if err != nil {
return fmt.Errorf("failed to install uv: %w\nPlease install uv v%s yourself: https://docs.astral.sh/uv/getting-started/installation/", err, UvVersion)
return fmt.Errorf("failed to install uv: %w", err)
}

_, _ = output.Write([]byte("\n"))
Expand All @@ -113,7 +126,7 @@ func (u *UvChecker) installUvCommand(ctx context.Context) error {
}

type uvInstaller interface {
EnsureUvInstalled(ctx context.Context) error
EnsureUvInstalled(ctx context.Context) (string, error)
}

type connectionFetcher interface {
Expand All @@ -125,17 +138,20 @@ type pipelineConnection interface {
}

type UvPythonRunner struct {
Cmd cmd
UvInstaller uvInstaller
conn connectionFetcher
Cmd cmd
UvInstaller uvInstaller
conn connectionFetcher
binaryFullPath string
}

func (u *UvPythonRunner) Run(ctx context.Context, execCtx *executionContext) error {
err := u.UvInstaller.EnsureUvInstalled(ctx)
binaryFullPath, err := u.UvInstaller.EnsureUvInstalled(ctx)
if err != nil {
return err
}

u.binaryFullPath = binaryFullPath

pythonVersion := "3.11"
if execCtx.asset.Image != "" {
image := execCtx.asset.Image
Expand All @@ -153,14 +169,15 @@ func (u *UvPythonRunner) Run(ctx context.Context, execCtx *executionContext) err
}

func (u *UvPythonRunner) RunIngestr(ctx context.Context, args []string, repo *git.Repo) error {
err := u.UvInstaller.EnsureUvInstalled(ctx)
binaryFullPath, err := u.UvInstaller.EnsureUvInstalled(ctx)
if err != nil {
return err
}
u.binaryFullPath = binaryFullPath

ingestrPackageName := "ingestr@" + ingestrVersion
err = u.Cmd.Run(ctx, repo, &command{
Name: "uv",
Name: u.binaryFullPath,
Args: []string{"tool", "install", "--force", "--quiet", "--python", pythonVersionForIngestr, ingestrPackageName},
})
if err != nil {
Expand All @@ -171,7 +188,7 @@ func (u *UvPythonRunner) RunIngestr(ctx context.Context, args []string, repo *gi
flags = append(flags, args...)

noDependencyCommand := &command{
Name: "uv",
Name: u.binaryFullPath,
Args: flags,
EnvVars: map[string]string{},
}
Expand All @@ -188,7 +205,7 @@ func (u *UvPythonRunner) runWithNoMaterialization(ctx context.Context, execCtx *
flags = append(flags, "--module", execCtx.module)

noDependencyCommand := &command{
Name: "uv",
Name: u.binaryFullPath,
Args: flags,
EnvVars: execCtx.envVariables,
}
Expand Down Expand Up @@ -236,7 +253,7 @@ func (u *UvPythonRunner) runWithMaterialization(ctx context.Context, execCtx *ex
flags = append(flags, tempPyScript.Name())

err = u.Cmd.Run(ctx, execCtx.repo, &command{
Name: "uv",
Name: u.binaryFullPath,
Args: flags,
EnvVars: execCtx.envVariables,
})
Expand Down Expand Up @@ -320,7 +337,7 @@ func (u *UvPythonRunner) runWithMaterialization(ctx context.Context, execCtx *ex

ingestrPackageName := "ingestr@" + ingestrVersion
err = u.Cmd.Run(ctx, execCtx.repo, &command{
Name: "uv",
Name: u.binaryFullPath,
Args: []string{"tool", "install", "--quiet", "--python", pythonVersionForIngestr, ingestrPackageName},
})
if err != nil {
Expand All @@ -337,7 +354,7 @@ func (u *UvPythonRunner) runWithMaterialization(ctx context.Context, execCtx *ex
}

err = u.Cmd.Run(ctx, execCtx.repo, &command{
Name: "uv",
Name: u.binaryFullPath,
Args: runArgs,
})
if err != nil {
Expand Down
16 changes: 8 additions & 8 deletions pkg/python/uv_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,9 @@ type mockUvInstaller struct {
mock.Mock
}

func (m *mockUvInstaller) EnsureUvInstalled(ctx context.Context) error {
func (m *mockUvInstaller) EnsureUvInstalled(ctx context.Context) (string, error) {
called := m.Called(ctx)
return called.Error(0)
return called.String(0), called.Error(1)
}

func Test_uvPythonRunner_Run(t *testing.T) {
Expand All @@ -41,12 +41,12 @@ func Test_uvPythonRunner_Run(t *testing.T) {
fields: func() *fields {
cmd := new(mockCmd)
cmd.On("Run", mock.Anything, repo, &command{
Name: "uv",
Name: "~/.bruin/uv",
Args: []string{"run", "--python", "3.11", "--module", module},
}).Return(assert.AnError)

inst := new(mockUvInstaller)
inst.On("EnsureUvInstalled", mock.Anything).Return(nil)
inst.On("EnsureUvInstalled", mock.Anything).Return("~/.bruin/uv", nil)

return &fields{
cmd: cmd,
Expand All @@ -68,12 +68,12 @@ func Test_uvPythonRunner_Run(t *testing.T) {
fields: func() *fields {
cmd := new(mockCmd)
cmd.On("Run", mock.Anything, repo, &command{
Name: "uv",
Name: "~/.bruin/uv",
Args: []string{"run", "--python", "3.11", "--with-requirements", "/path/to/requirements.txt", "--module", module},
}).Return(assert.AnError)

inst := new(mockUvInstaller)
inst.On("EnsureUvInstalled", mock.Anything).Return(nil)
inst.On("EnsureUvInstalled", mock.Anything).Return("~/.bruin/uv", nil)

return &fields{
cmd: cmd,
Expand All @@ -95,12 +95,12 @@ func Test_uvPythonRunner_Run(t *testing.T) {
fields: func() *fields {
cmd := new(mockCmd)
cmd.On("Run", mock.Anything, repo, &command{
Name: "uv",
Name: "~/.bruin/uv",
Args: []string{"run", "--python", "3.13", "--with-requirements", "/path/to/requirements.txt", "--module", module},
}).Return(assert.AnError)

inst := new(mockUvInstaller)
inst.On("EnsureUvInstalled", mock.Anything).Return(nil)
inst.On("EnsureUvInstalled", mock.Anything).Return("~/.bruin/uv", nil)

return &fields{
cmd: cmd,
Expand Down
12 changes: 12 additions & 0 deletions pkg/user/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,18 @@ func (c *ConfigManager) RecreateHomeDir() error {
return nil
}

func (c *ConfigManager) EnsureAndGetBruinHomeDir() (string, error) {
err := c.EnsureHomeDirExists()
if err != nil {
return "", err
}
absPath, err := filepath.Abs(c.bruinHomeDir)
if err != nil {
return "", err
}
return absPath, nil
}

func (c *ConfigManager) EnsureHomeDirExists() error {
err := c.ensureHomeDirSet()
if err != nil {
Expand Down

0 comments on commit f729c14

Please sign in to comment.