diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml index 3effa2b..a904175 100644 --- a/.github/workflows/build-and-test.yml +++ b/.github/workflows/build-and-test.yml @@ -144,8 +144,6 @@ jobs: CIBW_ARCHS: ${{ matrix.arch }} CIBW_SKIP: cp314* cp38-* CIBW_MANYLINUX: manylinux_2_39 - CIBW_MUSHLINUX_X86_64: manylinux_2_39 - CIBW_MUSHLINUX_AARCH64: manylinux_2_39 CIBW_BEFORE_ALL_WINDOWS: | choco install llvm cmake ninja export IS_MUSL=0 diff --git a/README.md b/README.md index 483f3f9..e07fbee 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,67 @@ pgembed.get_server("/path/to/my/data/dir") PostgreSQL binaries are available at `pgembed.POSTGRES_BIN_PATH` if you need direct access to tools like `initdb`, `pg_ctl`, `psql`, or `pg_config`. +## Extensions + +pgembed supports optional PostgreSQL extensions as separate packages. Install the extensions you need: + +```bash +# Base installation (PostgreSQL only) +pip install pgembed + +# With specific extensions (separate wheels) +pip install pgembed-pgvector +pip install pgembed-pgvectorscale +pip install pgembed-pgtextsearch + +# Multiple extensions +pip install pgembed-pgvector pgembed-pgvectorscale pgembed-pgtextsearch +``` + +Available extensions: +- `pgembed-pgvector`: Vector similarity search (works on all platforms) +- `pgembed-pgvectorscale`: High-performance vector storage (requires Rust, not available on Alpine/Windows) +- `pgembed-pgtextsearch`: BM25-based full-text search (requires Rust, not available on Alpine/Windows) + +### Checking available extensions + +```python +import pgembed + +# Check which extensions are available +print(pgembed.list_extensions()) +# {'pgvector': True, 'pgvectorscale': True, 'pgtextsearch': False, 'pg_duckdb': True} + +# Check if a specific extension is available +if pgembed.has_extension('pgvector'): + # Create the extension + server.create_extension('vector') +``` + +### Platform Support + +- **pgvector**: Works on Linux, macOS (Intel & Apple Silicon), Windows +- **pgvectorscale**: Works on Linux, macOS (Intel & Apple Silicon). NOT available on Alpine Linux or Windows (requires Rust) +- **pgtextsearch**: Works on Linux, macOS (Intel & Apple Silicon). NOT available on Alpine Linux or Windows (requires Rust) + +### Building specific extensions + +To build only specific extensions: + +```bash +# Build only pgvector +make pgvector + +# Build only pgvectorscale +make pgvectorscale + +# Build only pgtextsearch +make pgtextsearch + +# Build specific combination +make EXTENSIONS="pgvector pgtextsearch" all +``` + ## History pgembed is a fork of [pgserver](https://github.com/orm011/pgserver), which was inspired by [postgresql-wheel](https://github.com/michelp/postgresql-wheel). While those projects focused primarily on Linux wheels, pgembed extends the approach with: diff --git a/pgbuild/Makefile b/pgbuild/Makefile index 339ae79..c526631 100644 --- a/pgbuild/Makefile +++ b/pgbuild/Makefile @@ -18,15 +18,24 @@ else RUST_EXTENSIONS_SUPPORTED := 1 endif -.PHONY: all -all: pgvector postgres -ifneq ($(OS),Windows_NT) -all: pg_duckdb -endif ifeq ($(RUST_EXTENSIONS_SUPPORTED),1) -all: pgtextsearch pgvectorscale +EXTENSIONS ?= pgvector pgvectorscale pgtextsearch +else +EXTENSIONS ?= pgvector +endif +ifneq ($(OS),Windows_NT) +EXTENSIONS += pg_duckdb endif +.PHONY: all +all: postgres $(EXTENSIONS) + +.PHONY: pgvector pgvectorscale pgtextsearch pg_duckdb +pgvector: postgres $(INSTALL_PREFIX)/lib/vector.so +pgvectorscale: postgres $(INSTALL_PREFIX)/lib/vectorscale.so +pgtextsearch: postgres $(INSTALL_PREFIX)/lib/pg_textsearch.so +pg_duckdb: postgres $(INSTALL_PREFIX)/lib/pg_duckdb.so + ### postgres POSTGRES_VERSION := 17-stable POSTGRES_SRC := REL_17_STABLE @@ -76,9 +85,6 @@ $(INSTALL_PREFIX)/lib/vector.so: $(PGVECTOR_DIR)/Makefile $(INSTALL_PREFIX)/bin/ && $(MAKE) -C $(PGVECTOR_DIR) PG_CONFIG=$(INSTALL_PREFIX)/bin/pg_config$(EXE_EXT) -j \ && $(MAKE) -C $(PGVECTOR_DIR) PG_CONFIG=$(INSTALL_PREFIX)/bin/pg_config$(EXE_EXT) install -.PHONY: pgvector -pgvector: postgres $(INSTALL_PREFIX)/lib/vector.so - ### pg_duckdb PG_DUCKDB_TAG := v1.1.1 PG_DUCKDB_REPO := https://github.com/duckdb/pg_duckdb @@ -97,9 +103,6 @@ else $(error pg_duckdb build is not supported on Windows due to cargo-pgrx issues) endif -.PHONY: pg_duckdb -pg_duckdb: postgres $(INSTALL_PREFIX)/lib/pg_duckdb.so - ### pgvectorscale PGVECTORSCALE_TAG := 0.5.1 PGVECTORSCALE_URL := https://github.com/timescale/pgvectorscale/archive/refs/tags/$(PGVECTORSCALE_TAG).tar.gz @@ -126,9 +129,6 @@ else $(error pgvectorscale build is not supported on musl libc distributions (e.g., Alpine) due to cargo-pgrx issues) endif -.PHONY: pgvectorscale -pgvectorscale: postgres $(INSTALL_PREFIX)/lib/vectorscale.so - ### pgtextsearch PGTEXTSEARCH_TAG := v0.5.0 PGTEXTSEARCH_URL := https://github.com/timescale/pg_textsearch/archive/refs/tags/$(PGTEXTSEARCH_TAG).tar.gz @@ -153,9 +153,6 @@ else $(error pgtextsearch build is not supported on musl libc distributions (e.g., Alpine) due to cargo-pgrx issues) endif -.PHONY: pgtextsearch -pgtextsearch: postgres $(INSTALL_PREFIX)/lib/pg_textsearch.so - ### other .PHONY: clean clean-all clean: diff --git a/pyproject.toml b/pyproject.toml index 5e584b4..b977b68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ test = [ [tool.setuptools.packages.find] where = ["src"] # list of folders that contain the packages (["."] by default) -include = ["pgembed*"] # package names should match these glob patterns (["*"] by default) +include = ["pgembed*", "pgembed_pg*"] # package names should match these glob patterns [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/src/pgembed/__init__.py b/src/pgembed/__init__.py index e7e95b4..e31fd92 100644 --- a/src/pgembed/__init__.py +++ b/src/pgembed/__init__.py @@ -1,2 +1,121 @@ from ._commands import * from .postgres_server import PostgresServer, get_server +from pathlib import Path +from typing import Optional +import logging +import importlib.util + +_logger = logging.getLogger('pgembed') + +def _get_pkg_path(): + spec = importlib.util.find_spec('pgembed') + if spec and spec.submodule_search_locations: + return Path(spec.submodule_search_locations[0]) + return Path(__file__).parent + +EXTENSION_LIB_PATH = _get_pkg_path() / "pginstall" / "lib" +EXTENSION_POSTGRES_LIB_PATH = EXTENSION_LIB_PATH / "postgresql" + +AVAILABLE_EXTENSIONS = {} + +EXTENSION_PACKAGES = { + 'pgvector': 'pgembed_pgvector', + 'pgvectorscale': 'pgembed_pgvectorscale', + 'pgtextsearch': 'pgembed_pgtextsearch', +} + +EXTENSION_SO_FILES = { + 'pgvector': 'vector.so', + 'pgvectorscale': 'vectorscale-0.5.1.so', + 'pgtextsearch': 'pg_textsearch.so', + 'pg_duckdb': 'pg_duckdb.so', +} + +def _detect_extensions(): + global AVAILABLE_EXTENSIONS + + for name, pkg_name in EXTENSION_PACKAGES.items(): + try: + ext_pkg = __import__(pkg_name) + ext_path = ext_pkg.get_extension_path() + if ext_path and ext_path.exists(): + AVAILABLE_EXTENSIONS[name] = True + _logger.info(f"Detected extension from package {pkg_name}: {name}") + continue + except ImportError: + pass + + so_file = EXTENSION_SO_FILES.get(name) + if so_file: + bundled_path = EXTENSION_POSTGRES_LIB_PATH / so_file + if bundled_path.exists(): + AVAILABLE_EXTENSIONS[name] = True + _logger.info(f"Detected extension from bundled lib: {name}") + continue + + AVAILABLE_EXTENSIONS[name] = False + +def has_extension(name: str) -> bool: + """Check if a specific extension is available. + + Args: + name: Extension name (e.g., 'pgvector', 'pgvectorscale', 'pgtextsearch', 'pg_duckdb') + + Returns: + True if the extension is available, False otherwise. + """ + return AVAILABLE_EXTENSIONS.get(name, False) + +def list_extensions() -> dict: + """Return a dictionary of available extensions. + + Returns: + Dict mapping extension names to availability (True/False) + """ + return AVAILABLE_EXTENSIONS.copy() + +def get_extension_create_name(name: str) -> str: + """Get the SQL extension creation name for an extension. + + Args: + name: Extension name (e.g., 'pgvector', 'pgtextsearch') + + Returns: + The SQL name to use when creating the extension. + """ + create_names = { + 'pgvector': 'vector', + 'pgvectorscale': 'vectorscale', + 'pgtextsearch': 'pg_textsearch', + 'pg_duckdb': 'pg_duckdb', + } + return create_names.get(name, name) + +def get_extension_path(name: str) -> Optional[Path]: + """Get the path to an extension .so file. + + Args: + name: Extension name (e.g., 'pgvector', 'pgtextsearch') + + Returns: + Path to the .so file, or None if not available. + """ + pkg_name = EXTENSION_PACKAGES.get(name) + if pkg_name: + try: + ext_pkg = __import__(pkg_name) + ext_path = ext_pkg.get_extension_path() + if ext_path: + return ext_path + except ImportError: + pass + + so_file = EXTENSION_SO_FILES.get(name) + if so_file: + bundled_path = EXTENSION_POSTGRES_LIB_PATH / so_file + if bundled_path.exists(): + return bundled_path + + return None + +_detect_extensions() diff --git a/src/pgembed/_commands.py b/src/pgembed/_commands.py index 5db8984..2691bf0 100644 --- a/src/pgembed/_commands.py +++ b/src/pgembed/_commands.py @@ -4,8 +4,27 @@ from typing import Optional, List, Callable import logging import tempfile +import importlib.util -POSTGRES_BIN_PATH = Path(__file__).parent / "pginstall" / "bin" +def _get_pkg_path(): + spec = importlib.util.find_spec('pgembed') + if spec and spec.submodule_search_locations: + return Path(spec.submodule_search_locations[0]) + return Path(__file__).parent + +_pkg_path = _get_pkg_path() +POSTGRES_BIN_PATH = _pkg_path / "pginstall" / "bin" + +_postgres_binaries_available = POSTGRES_BIN_PATH.exists() + +if not _postgres_binaries_available: + import os + _logger = logging.getLogger('pgembed') + _logger.warning( + f"PostgreSQL binaries not found at {POSTGRES_BIN_PATH}. " + f"This is expected during development with editable install. " + f"Run 'make build' to build PostgreSQL binaries." + ) _logger = logging.getLogger('pgembed') @@ -60,6 +79,8 @@ def command(args : List[str], pgdata : Optional[Path] = None, **kwargs) -> str: __all__ = [] def _init(): + if not _postgres_binaries_available: + return for path in POSTGRES_BIN_PATH.iterdir(): exe_name = path.name prog = create_command_function(exe_name) diff --git a/src/pgembed/postgres_server.py b/src/pgembed/postgres_server.py index 4d1793c..5834793 100644 --- a/src/pgembed/postgres_server.py +++ b/src/pgembed/postgres_server.py @@ -9,7 +9,8 @@ import psutil import time -from ._commands import POSTGRES_BIN_PATH, initdb, pg_ctl +from . import _commands +from ._commands import POSTGRES_BIN_PATH from .utils import find_suitable_port, find_suitable_socket_dir, DiskList, PostmasterInfo, process_is_running if platform.system() != 'Windows': @@ -17,6 +18,16 @@ _logger = logging.getLogger('pgembed') +def _get_command(name: str): + cmd = getattr(_commands, name, None) + if cmd is None: + raise RuntimeError( + f"PostgreSQL binaries not available. " + f"pgembed was installed without PostgreSQL binaries or they were not built. " + f"Run 'make build' in the pgembed source directory." + ) + return cmd + class PostgresServer: """ Provides a common interface for interacting with a server. """ @@ -125,6 +136,7 @@ def ensure_pgdata_inited(self) -> None: proc.kill() assert not proc.is_running() + initdb = _get_command('initdb') initdb(['--auth=trust', '--auth-local=trust', '--encoding=utf8', '-U', self.postgres_user], pgdata=self.pgdata, user=self.system_user) else: @@ -173,6 +185,7 @@ def ensure_postgres_running(self) -> None: try: _logger.info(f"running pg_ctl... {pg_ctl_args=}") + pg_ctl = _get_command('pg_ctl') pg_ctl(pg_ctl_args,pgdata=self.pgdata, user=self.system_user, timeout=10) except subprocess.CalledProcessError as err: _logger.error(f"Failed to start server.\nShowing contents of postgres server log ({self.log.absolute()}) below:\n{self.log.read_text()}") @@ -217,6 +230,7 @@ def _cleanup(self) -> None: if self._postmaster_info is not None: if self._postmaster_info.process.is_running(): try: + pg_ctl = _get_command('pg_ctl') pg_ctl(['-w', 'stop'], pgdata=self.pgdata, user=self.system_user) stopped = True except subprocess.CalledProcessError: @@ -248,6 +262,36 @@ def psql(self, command : str) -> str: input=command.encode(), shell=True) return stdout.decode("utf-8") + def create_extension(self, extension_name: str) -> str: + """Create a PostgreSQL extension. + + Args: + extension_name: Name of the extension to create (e.g., 'vector', 'pg_textsearch') + + Returns: + The output from the CREATE EXTENSION command. + + Raises: + RuntimeError: If the extension is not available in the installation. + """ + import pgembed + from pgembed import AVAILABLE_EXTENSIONS, get_extension_create_name + extension_map = { + 'vector': 'pgvector', + 'vectorscale': 'pgvectorscale', + 'pg_textsearch': 'pgtextsearch', + 'pg_duckdb': 'pg_duckdb', + } + + pkg_name = extension_map.get(extension_name, extension_name) + if not pgembed.has_extension(pkg_name): + raise RuntimeError( + f"Extension '{extension_name}' is not available. " + f"Available extensions: {[k for k, v in AVAILABLE_EXTENSIONS.items() if v]}" + ) + create_name = get_extension_create_name(pkg_name) + return self.psql(f"CREATE EXTENSION IF NOT EXISTS {create_name};") + def __enter__(self): self._count += 1 return self diff --git a/src/pgembed_pgtextsearch/__init__.py b/src/pgembed_pgtextsearch/__init__.py new file mode 100644 index 0000000..9e30c5d --- /dev/null +++ b/src/pgembed_pgtextsearch/__init__.py @@ -0,0 +1,40 @@ +__version__ = "0.1.8" + +EXTENSION_NAME = "pgtextsearch" +EXTENSION_SO = "pg_textsearch.so" +EXTENSION_CREATE = "pg_textsearch" + +def get_extension_path(): + from pathlib import Path + pkg_dir = Path(__file__).parent + so_path = pkg_dir / EXTENSION_SO + if so_path.exists(): + return so_path + + try: + import pgembed + base_lib = pgembed.EXTENSION_LIB_PATH + bundled = base_lib / "postgresql" / EXTENSION_SO + if bundled.exists(): + return bundled + except ImportError: + pass + + return None + +def get_extension_share_path(): + from pathlib import Path + try: + import pgembed + base_share = Path(__file__).parent / "pginstall" / "share" / "postgresql" / "extension" + control_file = base_share / f"{EXTENSION_NAME}.control" + if control_file.exists(): + return base_share + + base_share = pgembed.EXTENSION_LIB_PATH.parent / "share" / "postgresql" / "extension" + control_file = base_share / f"{EXTENSION_NAME}.control" + if control_file.exists(): + return base_share + except ImportError: + pass + return None diff --git a/src/pgembed_pgtextsearch/pyproject.toml b/src/pgembed_pgtextsearch/pyproject.toml new file mode 100644 index 0000000..b0c951b --- /dev/null +++ b/src/pgembed_pgtextsearch/pyproject.toml @@ -0,0 +1,17 @@ +[project] +name = "pgembed-pgtextsearch" +version = "0.1.8" +description = "pg_textsearch extension for pgembed" +readme = "README.md" +requires-python = ">=3.12" +dependencies = ["pgembed>=0.1.8"] + +[project.optional-dependencies] +test = ["pytest"] + +[tool.setuptools.packages.find] +where = ["src"] + +[build-system] +requires = ["setuptools>=58.0.0", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/src/pgembed_pgvector/__init__.py b/src/pgembed_pgvector/__init__.py new file mode 100644 index 0000000..e32bca9 --- /dev/null +++ b/src/pgembed_pgvector/__init__.py @@ -0,0 +1,40 @@ +__version__ = "0.1.8" + +EXTENSION_NAME = "pgvector" +EXTENSION_SO = "vector.so" +EXTENSION_CREATE = "vector" + +def get_extension_path(): + from pathlib import Path + pkg_dir = Path(__file__).parent + so_path = pkg_dir / EXTENSION_SO + if so_path.exists(): + return so_path + + try: + import pgembed + base_lib = pgembed.EXTENSION_LIB_PATH + bundled = base_lib / "postgresql" / EXTENSION_SO + if bundled.exists(): + return bundled + except ImportError: + pass + + return None + +def get_extension_share_path(): + from pathlib import Path + try: + import pgembed + base_share = Path(__file__).parent / "pginstall" / "share" / "postgresql" / "extension" + control_file = base_share / f"{EXTENSION_NAME}.control" + if control_file.exists(): + return base_share + + base_share = pgembed.EXTENSION_LIB_PATH.parent / "share" / "postgresql" / "extension" + control_file = base_share / f"{EXTENSION_NAME}.control" + if control_file.exists(): + return base_share + except ImportError: + pass + return None diff --git a/src/pgembed_pgvector/pyproject.toml b/src/pgembed_pgvector/pyproject.toml new file mode 100644 index 0000000..226af69 --- /dev/null +++ b/src/pgembed_pgvector/pyproject.toml @@ -0,0 +1,17 @@ +[project] +name = "pgembed-pgvector" +version = "0.1.8" +description = "pgvector extension for pgembed" +readme = "README.md" +requires-python = ">=3.12" +dependencies = ["pgembed>=0.1.8"] + +[project.optional-dependencies] +test = ["pytest"] + +[tool.setuptools.packages.find] +where = ["src"] + +[build-system] +requires = ["setuptools>=58.0.0", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/src/pgembed_pgvectorscale/__init__.py b/src/pgembed_pgvectorscale/__init__.py new file mode 100644 index 0000000..964cc3a --- /dev/null +++ b/src/pgembed_pgvectorscale/__init__.py @@ -0,0 +1,40 @@ +__version__ = "0.1.8" + +EXTENSION_NAME = "pgvectorscale" +EXTENSION_SO = "vectorscale-0.5.1.so" +EXTENSION_CREATE = "vectorscale" + +def get_extension_path(): + from pathlib import Path + pkg_dir = Path(__file__).parent + so_path = pkg_dir / EXTENSION_SO + if so_path.exists(): + return so_path + + try: + import pgembed + base_lib = pgembed.EXTENSION_LIB_PATH + bundled = base_lib / "postgresql" / EXTENSION_SO + if bundled.exists(): + return bundled + except ImportError: + pass + + return None + +def get_extension_share_path(): + from pathlib import Path + try: + import pgembed + base_share = Path(__file__).parent / "pginstall" / "share" / "postgresql" / "extension" + control_file = base_share / f"{EXTENSION_NAME}.control" + if control_file.exists(): + return base_share + + base_share = pgembed.EXTENSION_LIB_PATH.parent / "share" / "postgresql" / "extension" + control_file = base_share / f"{EXTENSION_NAME}.control" + if control_file.exists(): + return base_share + except ImportError: + pass + return None diff --git a/src/pgembed_pgvectorscale/pyproject.toml b/src/pgembed_pgvectorscale/pyproject.toml new file mode 100644 index 0000000..2f9f0fc --- /dev/null +++ b/src/pgembed_pgvectorscale/pyproject.toml @@ -0,0 +1,17 @@ +[project] +name = "pgembed-pgvectorscale" +version = "0.1.8" +description = "pgvectorscale extension for pgembed" +readme = "README.md" +requires-python = ">=3.12" +dependencies = ["pgembed>=0.1.8"] + +[project.optional-dependencies] +test = ["pytest"] + +[tool.setuptools.packages.find] +where = ["src"] + +[build-system] +requires = ["setuptools>=58.0.0", "wheel"] +build-backend = "setuptools.build_meta"