diff --git a/memori/__init__.py b/memori/__init__.py index 3b15f627..b76bb7ff 100644 --- a/memori/__init__.py +++ b/memori/__init__.py @@ -13,7 +13,7 @@ from typing import Any from uuid import uuid4 -import psycopg +# Import psycopg lazily in _get_default_connection to avoid requiring it at import time from memori._config import Config from memori._exceptions import ( @@ -94,6 +94,13 @@ def __init__(self, conn: Callable[[], Any] | Any | None = None): def _get_default_connection(self) -> Callable[[], Any]: connection_string = os.environ.get("MEMORI_COCKROACHDB_CONNECTION_STRING") if connection_string: + try: + import psycopg + except ImportError as e: + raise RuntimeError( + "psycopg is required to use the default connection. Install 'psycopg' or pass a connection factory via 'conn'." + ) from e + return lambda: psycopg.connect(connection_string) raise RuntimeError( diff --git a/memori/storage/cockroachdb/_display.py b/memori/storage/cockroachdb/_display.py index 2d2a7ca0..392a3383 100644 --- a/memori/storage/cockroachdb/_display.py +++ b/memori/storage/cockroachdb/_display.py @@ -1,6 +1,7 @@ r""" __ __ _ | \/ | ___ _ __ ___ ___ _ __(_) + | |\/| |/ _ \ '_ ` _ \ / _ \| '__| | | | | | __/ | | | | | (_) | | | | |_| |_|\___|_| |_| |_|\___/|_| |_| @@ -12,8 +13,41 @@ class Display: - def __init__(self): - self.files = Files() + def __init__(self, files: Files | None = None, colorize: bool = False): + """Display helper for CockroachDB commands. + + Parameters + ---------- + files: optional Files + Optional `Files` instance to use (injected for testing). If not + provided, a default `Files()` will be created. + colorize: bool + When True, ASCII color/formatting sequences will be included in + the returned strings (useful for terminal output). Defaults to + False for predictable programmatic behavior. + """ + self.files = files or Files() + self.colorize = bool(colorize) + + def _color(self, text: str, color: str) -> str: + """Return text wrapped in simple ANSI color codes when colorize is True. + + Supported colors: 'green', 'yellow', 'blue', 'bold'. If an unknown + color is requested the text is returned unchanged. + """ + if not self.colorize: + return text + + codes = { + "green": "\x1b[32m", + "yellow": "\x1b[33m", + "blue": "\x1b[34m", + "bold": "\x1b[1m", + "reset": "\x1b[0m", + } + + start = codes.get(color, "") + return f"{start}{text}{codes['reset']}" if start else text def cluster_already_started(self): return """You already have an active CockroachDB cluster running. To start @@ -28,3 +62,48 @@ def cluster_was_not_started(self): python -m memori cockroachdb cluster start """ + + def banner(self): + """Return the ASCII art banner for CockroachDB helper messages.""" + return __doc__ + + def cluster_status(self): + """Return a user-friendly message describing the current cluster state. + + If a cluster id is present in `Files`, this will include the id and a + short hint to delete it. Otherwise, it returns the same message as + :py:meth:`cluster_was_not_started`. + """ + cid = self.files.read_id() + if cid: + cid_str = self._color(cid, "bold") + return f"""Your active CockroachDB cluster id is: {cid_str} +To delete it, run: + + python -m memori cockroachdb cluster delete +""" + + return self.cluster_was_not_started() + + def connection_string(self): + """Return a short example connection string for the active cluster. + + If no cluster is active, returns a hint that the user should start one. + """ + cid = self.files.read_id() + if cid: + host = self._color(cid, "green") + cmd = f"cockroach sql --insecure --host {host}:26257" + return cmd + + return "No active CockroachDB cluster found. Start one with:\n\n python -m memori cockroachdb cluster start\n" + + def example_connection_block(self): + """Return a short multi-line example for connecting to the active cluster.""" + cid = self.files.read_id() + if cid: + cmd = self.connection_string() + return f"To connect to your cluster, run:\n\n {cmd}\n" + + return "Start a cluster first to see a connection example.\n" + diff --git a/tests/storage/cockroachdb/test_storage_cockroachdb_display_color.py b/tests/storage/cockroachdb/test_storage_cockroachdb_display_color.py new file mode 100644 index 00000000..07507e3f --- /dev/null +++ b/tests/storage/cockroachdb/test_storage_cockroachdb_display_color.py @@ -0,0 +1,36 @@ +from memori.storage.cockroachdb._display import Display + + +class DummyFilesWithID: + def __init__(self, id="cluster-xyz"): + self._id = id + + def read_id(self): + return self._id + + +def test_colorize_connection_string_has_ansi(): + d = Display(files=DummyFilesWithID(), colorize=True) + cs = d.connection_string() + assert "cockroach sql" in cs + assert "\x1b[" in cs + + +def test_no_color_by_default(): + d = Display(files=DummyFilesWithID()) + cs = d.connection_string() + assert "cockroach sql" in cs + assert "\x1b[" not in cs + + +def test_example_block_when_started_colorized(): + d = Display(files=DummyFilesWithID(), colorize=True) + blk = d.example_connection_block() + assert "To connect to your cluster" in blk + assert "\x1b[" in blk + + +def test_example_block_when_not_started(): + d = Display(files=type("F", (), {"read_id": lambda self: None})()) + blk = d.example_connection_block() + assert "Start a cluster first" in blk diff --git a/tests/storage/cockroachdb/test_storage_cockroachdb_display_injection.py b/tests/storage/cockroachdb/test_storage_cockroachdb_display_injection.py new file mode 100644 index 00000000..2cb95efe --- /dev/null +++ b/tests/storage/cockroachdb/test_storage_cockroachdb_display_injection.py @@ -0,0 +1,12 @@ +from memori.storage.cockroachdb._display import Display + + +class DummyFiles: + def __init__(self): + self.called = True + + +def test_files_injection(): + dummy = DummyFiles() + d = Display(files=dummy) + assert d.files is dummy diff --git a/tests/storage/cockroachdb/test_storage_cockroachdb_display_status.py b/tests/storage/cockroachdb/test_storage_cockroachdb_display_status.py new file mode 100644 index 00000000..d4dbbcb1 --- /dev/null +++ b/tests/storage/cockroachdb/test_storage_cockroachdb_display_status.py @@ -0,0 +1,42 @@ +from memori.storage.cockroachdb._display import Display + + +class DummyFilesWithID: + def __init__(self, id="cluster-123"): + self._id = id + + def read_id(self): + return self._id + + +class DummyFilesNoID: + def read_id(self): + return None + + +def test_banner_contains_branding(): + assert "perfectam memoriam" in Display().banner() + + +def test_cluster_status_when_started(): + d = Display(files=DummyFilesWithID()) + out = d.cluster_status() + assert "active CockroachDB cluster id" in out + assert "cluster-123" in out + assert "cluster delete" in out + + +def test_cluster_status_when_not_started(): + d = Display(files=DummyFilesNoID()) + assert d.cluster_status() == d.cluster_was_not_started() + + +def test_connection_string_when_started(): + d = Display(files=DummyFilesWithID()) + assert "cockroach sql" in d.connection_string() + assert "cluster-123" in d.connection_string() + + +def test_connection_string_when_not_started(): + d = Display(files=DummyFilesNoID()) + assert "No active CockroachDB cluster found" in d.connection_string() diff --git a/tests/test_lazy_psycopg_import.py b/tests/test_lazy_psycopg_import.py new file mode 100644 index 00000000..77c0abbe --- /dev/null +++ b/tests/test_lazy_psycopg_import.py @@ -0,0 +1,43 @@ +import builtins +import importlib +import os +import pytest + + +def test_import_without_psycopg(monkeypatch): + """Ensure importing the package works even when psycopg is not installed.""" + orig_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "psycopg" or name.startswith("psycopg."): + raise ImportError("No module named 'psycopg'") + return orig_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + # Import should succeed because psycopg is only required when using default connection + importlib.invalidate_caches() + memori = importlib.import_module("memori") + + assert hasattr(memori, "Memori") + + +def test_default_connection_requires_psycopg(monkeypatch): + """If MEMORI_COCKROACHDB_CONNECTION_STRING is set, importing and instantiating Memori should raise a clear error when psycopg is missing.""" + orig_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "psycopg" or name.startswith("psycopg."): + raise ImportError("No module named 'psycopg'") + return orig_import(name, globals, locals, fromlist, level) + + monkeypatch.setenv("MEMORI_COCKROACHDB_CONNECTION_STRING", "postgresql://user:pass@localhost/db") + monkeypatch.setattr(builtins, "__import__", fake_import) + + importlib.invalidate_caches() + memori = importlib.import_module("memori") + + with pytest.raises(RuntimeError) as e: + memori.Memori() + + assert "psycopg is required" in str(e.value)