import asyncio import json import logging import os import platform import shutil import stat import subprocess import sys from pathlib import Path from typing import Optional import httpx logger = logging.getLogger(__name__) DEFAULT_DATA_DIR = Path(os.environ.get("HINDSIGHT_API_PG0_DATA_DIR", Path.home() / ".hindsight" / "pg_data")) DEFAULT_INSTALL_DIR = Path.home() / ".hindsight" / "bin" BINARY_NAME = "pg0" DEFAULT_PORT = 5555 DEFAULT_USERNAME = "hindsight" DEFAULT_PASSWORD = "hindsight" DEFAULT_DATABASE = "hindsight" def get_platform_binary_name() -> str: """Get the appropriate binary name for the current platform. Supported platforms: - macOS ARM64 (darwin-aarch64) - Linux x86_64 - Windows x86_64 """ system = platform.system().lower() machine = platform.machine().lower() # Normalize architecture names if machine in ("x86_64", "amd64"): arch = "x86_64" elif machine in ("arm64", "aarch64"): arch = "aarch64" else: raise RuntimeError( f"Embedded PostgreSQL is not supported on architecture: {machine}. " f"Supported architectures: x86_64/amd64 (Linux, Windows), aarch64/arm64 (macOS)" ) if system == "darwin" and arch == "aarch64": return "pg0-darwin-aarch64" elif system == "linux" and arch == "x86_64": return "pg0-linux-x86_64" elif system == "windows" and arch == "x86_64": return "pg0-windows-x86_64.exe" else: raise RuntimeError( f"Embedded PostgreSQL is not supported on {system}-{arch}. " f"Supported platforms: darwin-aarch64 (macOS ARM), linux-x86_64, windows-x86_64" ) def get_download_url( version: str = "latest", repo: str = "vectorize-io/pg0", ) -> str: """ """ # Check for direct URL override binary_name = get_platform_binary_name() if version == "latest": return f"https://github.com/{repo}/releases/latest/download/{binary_name}" else: return f"https://github.com/{repo}/releases/download/{version}/{binary_name}" class EmbeddedPostgres: """ Manages an embedded PostgreSQL server instance. This class handles: - Downloading and installing the embedded-postgres CLI - Starting/stopping the PostgreSQL server - Getting the connection URI Example: pg = EmbeddedPostgres(data_dir="~/.myapp/data") await pg.ensure_installed() await pg.start() uri = await pg.get_uri() # ... use uri with asyncpg ... await pg.stop() """ def __init__( self, data_dir: Optional[Path] = None, install_dir: Optional[Path] = None, version: str = "latest", port: int = DEFAULT_PORT, username: str = DEFAULT_USERNAME, password: str = DEFAULT_PASSWORD, database: str = DEFAULT_DATABASE, name: str = "hindsight", ): """ Initialize the embedded PostgreSQL manager. Args: data_dir: Directory to store PostgreSQL data. Defaults to ~/.hindsight/pg_data install_dir: Directory to install the CLI binary. Defaults to ~/.hindsight/bin version: Version of embedded-postgres to use. Defaults to "latest" port: Port to listen on. Defaults to 5555 username: Username for the database. Defaults to "hindsight" password: Password for the database. Defaults to "hindsight" database: Database name to create. Defaults to "hindsight" name: Instance name for pg0. Defaults to "hindsight" """ self.data_dir = Path(data_dir or DEFAULT_DATA_DIR).expanduser() self.install_dir = Path(install_dir or DEFAULT_INSTALL_DIR).expanduser() self.version = version self.port = port self.username = username self.password = password self.database = database self.name = name # Binary path binary_name = "pg0.exe" if platform.system() == "Windows" else "pg0" self.binary_path = self.install_dir / binary_name self._process: Optional[subprocess.Popen] = None def is_installed(self) -> bool: """Check if the embedded-postgres CLI is installed.""" return self.binary_path.exists() and os.access(self.binary_path, os.X_OK) async def ensure_installed(self) -> None: """ Ensure the embedded-postgres CLI is installed. Downloads and installs the binary if not already present. """ if self.is_installed(): return logger.info("Installing pg0 CLI...") # Create install directory self.install_dir.mkdir(parents=True, exist_ok=True) # Download the binary download_url = get_download_url(self.version) logger.info(f"Downloading from {download_url}") try: async with httpx.AsyncClient(follow_redirects=True, timeout=300.0) as client: response = await client.get(download_url) response.raise_for_status() # Write binary to disk with open(self.binary_path, "wb") as f: f.write(response.content) # Make executable on Unix if platform.system() != "Windows": st = os.stat(self.binary_path) os.chmod(self.binary_path, st.st_mode | stat.S_IEXEC) logger.info(f"Installed pg0 to {self.binary_path}") except httpx.HTTPError as e: raise RuntimeError(f"Failed to download pg0: {e}") from e def _run_command(self, *args: str, capture_output: bool = True) -> subprocess.CompletedProcess: """Run an embedded-postgres command synchronously.""" cmd = [str(self.binary_path), *args] return subprocess.run( cmd, capture_output=capture_output, text=True, ) async def _run_command_async(self, *args: str) -> tuple[int, str, str]: """Run an embedded-postgres command asynchronously.""" cmd = [str(self.binary_path), *args] process = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) stdout, stderr = await process.communicate() return process.returncode, stdout.decode(), stderr.decode() async def start(self) -> str: """ Start the PostgreSQL server. Returns: The connection URI for the started server. Raises: RuntimeError: If the server fails to start. """ if not self.is_installed(): raise RuntimeError("pg0 is not installed. Call ensure_installed() first.") # Create data directory self.data_dir.mkdir(parents=True, exist_ok=True) logger.info(f"Starting embedded PostgreSQL (name: {self.name}, data: {self.data_dir}, install: {self.install_dir}, port: {self.port})...") returncode, stdout, stderr = await self._run_command_async( "start", "--name", self.name, "--port", str(self.port), "--username", self.username, "--password", self.password, "--database", self.database, "--data-dir", self.data_dir.as_posix() ) if returncode != 0: raise RuntimeError(f"Failed to start PostgreSQL: {stderr}") logger.info("Embedded PostgreSQL started") # Get and return the URI return await self.get_uri() async def stop(self) -> None: """ Stop the PostgreSQL server. Raises: RuntimeError: If the server fails to stop. """ if not self.is_installed(): return logger.info(f"Stopping embedded PostgreSQL (name: {self.name})...") returncode, stdout, stderr = await self._run_command_async("stop", "--name", self.name) if returncode != 0: # Don't raise if server wasn't running if "not running" in stderr.lower(): return raise RuntimeError(f"Failed to stop PostgreSQL: {stderr}") logger.info("Embedded PostgreSQL stopped") async def _get_info(self) -> dict: """ Get info from pg0 using the `info -o json` command. Returns: Dictionary with 'running' (bool) and 'uri' (str) keys. Raises: RuntimeError: If unable to get info. """ if not self.is_installed(): raise RuntimeError("pg0 is not installed.") returncode, stdout, stderr = await self._run_command_async( "info", "--name", self.name, "-o", "json") if returncode != 0: raise RuntimeError(f"Failed to get PostgreSQL info: {stderr}") try: return json.loads(stdout.strip()) except json.JSONDecodeError as e: raise RuntimeError(f"Failed to parse pg0 info output: {e}") async def get_uri(self) -> str: """ Get the connection URI for the PostgreSQL server. Returns: PostgreSQL connection URI (e.g., postgresql://user:pass@localhost:5432/db) Raises: RuntimeError: If unable to get the URI or server is not running. """ info = await self._get_info() uri = info.get("uri") if not uri: raise RuntimeError("PostgreSQL server is not running or URI not available") return uri async def status(self) -> dict: """ Get the status of the PostgreSQL server. Returns: Dictionary with status information including 'running' boolean and 'uri'. """ if not self.is_installed(): return {"installed": False, "running": False} try: info = await self._get_info() return { "installed": True, "running": info.get("running", False), "uri": info.get("uri"), "data_dir": str(self.data_dir), "binary_path": str(self.binary_path), } except RuntimeError: return { "installed": True, "running": False, "data_dir": str(self.data_dir), "binary_path": str(self.binary_path), } async def is_running(self) -> bool: """Check if the PostgreSQL server is currently running.""" if not self.is_installed(): return False try: info = await self._get_info() return info.get("running", False) except RuntimeError: return False async def ensure_running(self) -> str: """ Ensure the PostgreSQL server is running. Installs if needed, starts if not running. Returns: The connection URI. """ await self.ensure_installed() if await self.is_running(): return await self.get_uri() return await self.start() def uninstall(self) -> None: """Remove the embedded-postgres binary.""" if self.binary_path.exists(): self.binary_path.unlink() logger.info(f"Removed {self.binary_path}") def clear_data(self) -> None: """Remove all PostgreSQL data (destructive!).""" if self.data_dir.exists(): shutil.rmtree(self.data_dir) logger.info(f"Removed data directory {self.data_dir}") # Convenience functions for simple usage _default_instance: Optional[EmbeddedPostgres] = None def get_embedded_postgres( data_dir: Optional[Path] = None, install_dir: Optional[Path] = None, ) -> EmbeddedPostgres: """ Get or create the default EmbeddedPostgres instance. Args: data_dir: Override default data directory install_dir: Override default install directory Returns: EmbeddedPostgres instance """ global _default_instance if _default_instance is None or data_dir or install_dir: _default_instance = EmbeddedPostgres( data_dir=data_dir, install_dir=install_dir, ) return _default_instance async def start_embedded_postgres( data_dir: Optional[Path] = None, ) -> str: """ Quick start function for embedded PostgreSQL. Downloads, installs, and starts PostgreSQL in one call. Args: data_dir: Directory to store PostgreSQL data Returns: Connection URI string Example: db_url = await start_embedded_postgres() conn = await asyncpg.connect(db_url) """ pg = get_embedded_postgres(data_dir=data_dir) return await pg.ensure_running() async def stop_embedded_postgres() -> None: """Stop the default embedded PostgreSQL instance.""" global _default_instance if _default_instance: await _default_instance.stop()