trove/trovedb/db.py

"""
Python API for accessing underlying SQlite3 File Database

A Sqlite3 database provides the basic underlying interface to the underlying
notes interface. The low level database is an extremely simple model using
similar ideas to git storage.
"""

import argparse
import sqlite3
import sys
from datetime import datetime, timezone
from pathlib import Path
from .trove import NODE_ROOT_ID

_SCHEMA = """
CREATE TABLE IF NOT EXISTS objects (
    id       INTEGER PRIMARY KEY,
    type     TEXT NOT NULL CHECK(type IN ('blob', 'tree')),
    data     BLOB,
    modified TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS metadata (
    id    INTEGER NOT NULL REFERENCES objects(id) ON DELETE CASCADE,
    key   TEXT NOT NULL,
    value BLOB,
    PRIMARY KEY (id, key)
);

CREATE TABLE IF NOT EXISTS labels
(
    label TEXT PRIMARY KEY,
    id INTEGER NOT NULL REFERENCES objects(id) ON DELETE CASCADE
); 
"""


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _initialize_db(con: sqlite3.Connection):
    con.executescript(_SCHEMA)
    con.commit()


class Sqlite3Trove:
    def __init__(self, con: sqlite3.Connection):
        self._con = con
        self._con.execute("PRAGMA foreign_keys = ON")
        self._con.row_factory = sqlite3.Row

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    @classmethod
    def open(cls, path: str | Path, create: bool = False) -> "Sqlite3Trove":
        """Open an existing Trove database."""
        p = Path(path)
        initialize = False
        if not p.exists():
            if not create:
                raise FileNotFoundError(f"Database not found: {p}")
            initialize = True
        con = sqlite3.connect(str(p))
        if initialize:
            con.executescript(_SCHEMA)
            con.commit()
        obj = cls(con)
        if initialize:
            obj.write_blob(b"", NODE_ROOT_ID)
        return obj

    def close(self):
        self._con.close()

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()

    # ------------------------------------------------------------------
    # CRUD operations
    # ------------------------------------------------------------------

    def read_object(self, object_id: int) -> bytes | None:
        """Return raw data for a blob object, or None if not found."""
        row = self._con.execute(
            "SELECT data, type FROM objects WHERE id = ?", (object_id,)
        ).fetchone()
        if row is None:
            return None
        return bytes(row["data"]) if row["data"] is not None else b""

    def read_metadata(self, object_id: int, key: str) -> bytes | None:
        """Return raw metadata value for (uuid, key), or None if not found."""
        row = self._con.execute(
            "SELECT value FROM metadata WHERE id = ? AND key = ?", (object_id, key)
        ).fetchone()
        if row is None:
            return None
        return bytes(row["value"]) if row["value"] is not None else b""

    def _write_object(self, data: bytes, dtype: str, object_id: int | None = None) -> int:
        """
        Insert or replace an object. Returns the id.
        If object_id is None, creates a new object with auto-assigned ID.
        If object_id is provided, updates or creates the object with that ID.
        """
        modified = _now()
        if object_id is None:
            cur = self._con.execute(
                "INSERT INTO objects (type, data, modified) VALUES (?, ?, ?)",
                (dtype, data, modified)
            )
            self._con.commit()
            return cur.lastrowid
        else:
            self._con.execute(
                "INSERT OR REPLACE INTO objects (id, type, data, modified) VALUES (?, ?, ?, ?)",
                (object_id, dtype, data, modified)
            )
            self._con.commit()
            return object_id

    def write_blob(self, data: bytes, existing_id: int | None = None) -> int:
        """
        Insert or replace a blob. Returns the id.
        Pass existing_id to update an existing object.
        """
        return self._write_object(data, "blob", existing_id)

    def delete_object(self, object_id: int) -> bool:
        """
        Delete a blob and all its metadata rows.
        Returns True if an object was deleted, False if id not found.
        Foreign key cascade handles the metadata rows.
        """
        cur = self._con.execute(
            "DELETE FROM objects WHERE id = ?", (object_id,)
        )
        self._con.commit()
        return cur.rowcount > 0

    def get_label(self, label: str) -> int | None:
        """
        Return the ID associated with a label, or None if not found.
        """
        row = self._con.execute(
            "SELECT id FROM labels WHERE label = ?", (label,)
        ).fetchone()
        if row is None:
            return None
        return row["id"]

    def set_label(self, label: str, object_id: int) -> None:
        """
        Set a label to point to an ID. Creates or updates the label.
        The ID must exist in the objects table.
        """
        self._con.execute(
            "INSERT OR REPLACE INTO labels (label, id) VALUES (?, ?)",
            (label, object_id),
        )
        self._con.commit()

    def delete_label(self, label: str) -> bool:
        """
        Delete a label. Returns True if a label was deleted, False if not found.
        """
        cur = self._con.execute(
            "DELETE FROM labels WHERE label = ?", (label,)
        )
        self._con.commit()
        return cur.rowcount > 0

    def list_labels(self) -> list[tuple[str, int]]:
        """
        Return all labels as a list of (label, id) tuples.
        """
        rows = self._con.execute(
            "SELECT label, id FROM labels ORDER BY label"
        ).fetchall()
        return [(row["label"], row["id"]) for row in rows]


def main():
    """Command-line interface for TroveDB."""
    parser = argparse.ArgumentParser(description="TroveDB command-line interface")
    parser.add_argument("database", help="Path to the database file")

    subparsers = parser.add_subparsers(dest="operation", required=True, help="Operation to perform")

    # Create database
    subparsers.add_parser("create", help="Create a new database")

    # Get blob
    get_parser = subparsers.add_parser("get", help="Get a blob object by ID")
    get_parser.add_argument("id", type=int, help="ID of the blob to retrieve")

    # Write blob
    write_parser = subparsers.add_parser("write", help="Write data to a blob")
    write_parser.add_argument("data", help="Data to write (as string, will be encoded as UTF-8)")
    write_parser.add_argument("--id", type=int, help="ID of existing blob to update (optional)")

    # Delete blob
    delete_parser = subparsers.add_parser("delete", help="Delete a blob by ID")
    delete_parser.add_argument("id", type=int, help="ID of the blob to delete")

    # Set label
    setlabel_parser = subparsers.add_parser("setlabel", help="Create or update a label to point to an ID")
    setlabel_parser.add_argument("label", help="Label name")
    setlabel_parser.add_argument("id", type=int, help="ID to associate with the label")

    # Remove label
    rmlabel_parser = subparsers.add_parser("rmlabel", help="Delete a label")
    rmlabel_parser.add_argument("label", help="Label name to delete")

    # List labels
    subparsers.add_parser("labels", help="List all labels")

    args = parser.parse_args()

    try:
        match args.operation:
            case "create":
                if Path(args.database).exists():
                    print(f"Database already exists: {args.database}", file=sys.stderr)
                    sys.exit(1)
                db = Sqlite3Trove.open(args.database, create=True)
                db.close()
                print(f"Database created: {args.database}")

            case "get":
                db = Sqlite3Trove.open(args.database)
                data = db.read_object(args.id)
                db.close()
                if data is None:
                    print(f"Blob not found: {args.id}")
                    sys.exit(1)
                sys.stdout.buffer.write(data)

            case "write":
                data_bytes = args.data.encode("utf-8")
                db = Sqlite3Trove.open(args.database)
                object_id = db.write_blob(data_bytes, args.id)
                db.close()
                print(object_id)

            case "delete":
                db = Sqlite3Trove.open(args.database)
                deleted = db.delete_object(args.id)
                db.close()
                if deleted:
                    print(f"Deleted blob: {args.id}")
                else:
                    print(f"Blob not found: {args.id}")
                    sys.exit(1)

            case "setlabel":
                db = Sqlite3Trove.open(args.database)
                db.set_label(args.label, args.id)
                db.close()
                print(f"Label '{args.label}' set to ID: {args.id}")

            case "rmlabel":
                db = Sqlite3Trove.open(args.database)
                deleted = db.delete_label(args.label)
                db.close()
                if deleted:
                    print(f"Deleted label: {args.label}")
                else:
                    print(f"Label not found: {args.label}")
                    sys.exit(1)

            case "labels":
                db = Sqlite3Trove.open(args.database)
                labels = db.list_labels()
                db.close()
                if labels:
                    for label, id in labels:
                        print(f"{label}: {id}")
                else:
                    print("No labels found.")

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()