Implement python package extras

.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

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:

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:

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"]

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()

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)