diff --git a/airbyte/_executors/python.py b/airbyte/_executors/python.py index c6578c0d..d712761e 100644 --- a/airbyte/_executors/python.py +++ b/airbyte/_executors/python.py @@ -1,9 +1,9 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. from __future__ import annotations +import os import shlex import subprocess -import sys from contextlib import suppress from pathlib import Path from shutil import rmtree @@ -17,6 +17,7 @@ from airbyte._util.meta import is_windows from airbyte._util.telemetry import EventState, log_install_state from airbyte._util.venv_util import get_bin_dir +from airbyte._util import pip_util, uv_util if TYPE_CHECKING: @@ -75,19 +76,6 @@ def interpreter_path(self) -> Path: suffix: Literal[".exe", ""] = ".exe" if is_windows() else "" return get_bin_dir(self._get_venv_path()) / ("python" + suffix) - def _run_subprocess_and_raise_on_failure(self, args: list[str]) -> None: - result = subprocess.run( - args, - check=False, - stderr=subprocess.PIPE, - ) - if result.returncode != 0: - raise exc.AirbyteSubprocessFailedError( - run_args=args, - exit_code=result.returncode, - log_text=result.stderr.decode("utf-8"), - ) - def uninstall(self) -> None: if self._get_venv_path().exists(): rmtree(str(self._get_venv_path())) @@ -106,18 +94,15 @@ def install(self) -> None: After installation, the installed version will be stored in self.reported_version. """ - self._run_subprocess_and_raise_on_failure( - [sys.executable, "-m", "venv", str(self._get_venv_path())] - ) - - pip_path = str(get_bin_dir(self._get_venv_path()) / "pip") + uv_util.create_venv(str(self._get_venv_path())) print( f"Installing '{self.name}' into virtual environment '{self._get_venv_path()!s}'.\n" - f"Running 'pip install {self.pip_url}'...\n" + f"Running 'uv pip install {self.pip_url}'...\n" ) try: - self._run_subprocess_and_raise_on_failure( - args=[pip_path, "install", *shlex.split(self.pip_url)] + uv_util.install_package( + venv_path=self._get_venv_path(), + pip_url=self.pip_url, ) except exc.AirbyteSubprocessFailedError as ex: # If the installation failed, remove the virtual environment diff --git a/airbyte/_util/uv_util.py b/airbyte/_util/uv_util.py new file mode 100644 index 00000000..2df87bc3 --- /dev/null +++ b/airbyte/_util/uv_util.py @@ -0,0 +1,102 @@ +"""Utility functions for UV virtual environment management.""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +from airbyte import exceptions as exc + + +def run_uv_command( + args: list[str], + env: dict | None = None, + raise_on_error: bool = True, +) -> subprocess.CompletedProcess: + """Run a UV command and return the result. + + Args: + args: Command arguments to pass to UV + env: Optional environment variables + raise_on_error: Whether to raise an exception on non-zero return code + + Returns: + CompletedProcess instance with command output + """ + result = subprocess.run( + ["uv", *args], + capture_output=True, + text=True, + env=env, + check=False, + ) + if raise_on_error and result.returncode != 0: + raise exc.AirbyteSubprocessFailedError( + run_args=args, + exit_code=result.returncode, + log_text=result.stderr.decode("utf-8"), + ) + return result + + +def create_venv(venv_path: Path, with_pip: bool = True) -> None: + """Create a virtual environment using UV. + + Args: + venv_path: Path where the virtual environment should be created + with_pip: Whether to include pip and other seed packages + """ + args = ["venv"] + if with_pip: + args.append("--seed") + args.append(str(venv_path)) + run_uv_command(args) + + +def install_package(venv_path: Path, pip_url: str) -> None: + """Install a package into a virtual environment using UV. + + Args: + venv_path: Path to the virtual environment + pip_url: Package specification (name, URL, or path) to install + """ + venv_env = get_venv_env(venv_path) + run_uv_command( + ["pip", "install", pip_url], + env=venv_env, + ) + + +def get_venv_bin_path(venv_path: Path) -> Path: + """Get the path of the executable 'bin' folder for a virtual env.""" + return venv_path / ("Scripts" if os.name == "nt" else "bin") + + +def get_venv_env(venv_path: Path) -> dict: + """Get environment variables for a virtual environment. + + Args: + venv_path: Path to the virtual environment + + Returns: + Dict of environment variables + """ + env = os.environ.copy() + env["VIRTUAL_ENV"] = str(venv_path) + env["PATH"] = f"{get_venv_bin_path(venv_path)}:{os.environ['PATH']}" + return env + + +def get_executable_path(venv_path: Path, executable: str) -> Path: + """Get the path to an executable in a virtual environment. + + Args: + venv_path: Path to the virtual environment + executable: Name of the executable + + Returns: + Path to the executable + """ + suffix = ".exe" if os.name == "nt" else "" + return get_venv_bin_path(venv_path) / f"{executable}{suffix}" diff --git a/poetry.lock b/poetry.lock index f85ff6f6..1be059d1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -4168,6 +4168,35 @@ files = [ {file = "uuid7-0.1.0.tar.gz", hash = "sha256:8c57aa32ee7456d3cc68c95c4530bc571646defac01895cfc73545449894a63c"}, ] +[[package]] +name = "uv" +version = "0.1.45" +description = "An extremely fast Python package installer and resolver, written in Rust." +optional = false +python-versions = ">=3.8" +groups = ["main"] +markers = "python_version <= \"3.11\"" +files = [ + {file = "uv-0.1.45-py3-none-linux_armv6l.whl", hash = "sha256:088af576fb0e0462cd5f718d03fb1a9f16ce5ae61fdb2a9d3ea938fc826cecc1"}, + {file = "uv-0.1.45-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:b94180009264f3f7ee74250f8e4f99c8cb0cb3633e3a9c9c66cdef3eb69be575"}, + {file = "uv-0.1.45-py3-none-macosx_11_0_arm64.whl", hash = "sha256:4e5d55f0f8b6ae416c72d78106e224c8e8338356da21ddebecc7b1723de80924"}, + {file = "uv-0.1.45-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:7fdb235aaf420fa8ac9009999b1654a23540f03e25c35094543c2f48d7c41aef"}, + {file = "uv-0.1.45-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:de81501c0b03160d0944906d1a713f108258360e20c58385974acb7253b56166"}, + {file = "uv-0.1.45-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:346aa2d0a4ad3c0c3f7852c1edf5e5a8e5d2ef34c7474e9089877291c2da979c"}, + {file = "uv-0.1.45-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:a601eed14d484d36d421e4208911a56aaf758ea6c385ef8edf8ad9f8ead57ce1"}, + {file = "uv-0.1.45-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2ca2d5a5e06c5f71c7b213e14fa59129e63b77de3ffbcf84ecc98d647d73a821"}, + {file = "uv-0.1.45-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:90b68c80dddebeca69b26a2af1e2e683804bcf2b5f22d107af03d9156d6218c6"}, + {file = "uv-0.1.45-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cd7f2f64fdded940342dc37234c11ae3508222c3c9b6b0eac5879dcd586010fa"}, + {file = "uv-0.1.45-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:a39141e179fea043151a165c9155031e7976b0e4b076c0c33a45b58a420134e0"}, + {file = "uv-0.1.45-py3-none-musllinux_1_1_armv7l.whl", hash = "sha256:68718add6ee2cef2816f9bf8a1dbf2d8cf63d98ddf45840f340029f65a49fd89"}, + {file = "uv-0.1.45-py3-none-musllinux_1_1_i686.whl", hash = "sha256:110e0f45ddb2fe832ce50b0308be90e5439e0c02d3ffe042feeb3f759811f31f"}, + {file = "uv-0.1.45-py3-none-musllinux_1_1_ppc64le.whl", hash = "sha256:0f6cfe885f109bacc055edd5df2c837616ae2238b9324a9d37835a96b204ab2f"}, + {file = "uv-0.1.45-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:87e77d25e8f358c0d5de1983497ee4cf4cea8fc73373d1ef1063533352db2f89"}, + {file = "uv-0.1.45-py3-none-win32.whl", hash = "sha256:ddb93620c9e01fa83573c2648df4bee3fa548ca940de51c8a2c3566a23a0c776"}, + {file = "uv-0.1.45-py3-none-win_amd64.whl", hash = "sha256:8e2eeea4eec0e09f7d67378152428b5308dba8b33990d045d7a31d19bf18ca1f"}, + {file = "uv-0.1.45.tar.gz", hash = "sha256:40fab956bc7af50dfa4bda14e5871528f57603eb9bf8595eb3144aace0ed8c47"}, +] + [[package]] name = "viztracer" version = "0.16.3" @@ -4336,4 +4365,4 @@ files = [ [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.12" -content-hash = "0e862e89b7b7b40c7aacacd5e9d277c62afc27dc8ff465ec3085e77aaad561ad" +content-hash = "21eba06a8d204e30e69f24342944a1ee59f4ddc48fde7be98309d2d1b326110e" diff --git a/pyproject.toml b/pyproject.toml index ca8a19af..fde59435 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ sqlalchemy = ">=1.4.51,!=2.0.36,<3.0" sqlalchemy-bigquery = { version = "1.12.0", python = "<3.13" } typing-extensions = "*" uuid7 = "^0.1.0" +uv = { version = "^0.1.45" } [tool.poetry.group.dev.dependencies] coverage = "^7.5.1" diff --git a/tests/integration_tests/test_uv_functionality.py b/tests/integration_tests/test_uv_functionality.py new file mode 100644 index 00000000..233a60b0 --- /dev/null +++ b/tests/integration_tests/test_uv_functionality.py @@ -0,0 +1,64 @@ +"""Tests for UV functionality and integration validation.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from airbyte._util.uv_util import ( + create_venv, + get_executable_path, + get_venv_env, + install_package, + run_uv_command, +) + + +def test_uv_available() -> None: + """Test that UV is available on the command line.""" + result = run_uv_command(["--version"]) + assert result.returncode == 0 + assert "uv" in result.stdout.lower() + + +@pytest.fixture +def temp_venv(tmp_path: Path) -> Path: + """Create a temporary directory for virtual environment testing. + + Uses pytest's tmp_path fixture which handles Windows cleanup properly. + """ + venv_path = tmp_path / ".venv" + return venv_path + + +def test_uv_venv_creation(temp_venv: Path) -> None: + """Test that UV can create a virtual environment.""" + result = run_uv_command(["venv", str(temp_venv)]) + assert result.returncode == 0 + assert temp_venv.exists() + assert (temp_venv / "bin" / "python").exists() + + +def test_uv_package_install(temp_venv: Path) -> None: + """Test that UV can install a package and we can execute it.""" + # Create venv and install black + create_venv(temp_venv, with_pip=True) + install_package( + temp_venv, "black" + ) # Use black since it's a common tool that creates executables + + # Verify package executable exists and is runnable + black_path = get_executable_path(temp_venv, "black") + assert black_path.exists() + + # Try running the package + result = subprocess.run( + [str(black_path), "--version"], + env=get_venv_env(temp_venv), + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "black" in result.stdout.lower()