fix: lazy psycopg import and improve CockroachDB display

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(

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

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

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

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

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)