cleveragents-core/features/steps/database_resources_steps.py

"""Step definitions for database resource types and transaction sandbox.

Covers:
- Database type definitions (postgres, mysql, sqlite, duckdb)
- Connection argument resolution and env var fallback
- Credential masking in error/validation messages
- Read-only enforcement
- Transaction sandbox create/commit/rollback lifecycle
- SQLite resource handler roundtrip
- DatabaseResourceHandler protocol conformance
- SandboxFactory transaction_rollback routing
"""

from __future__ import annotations

import os
import sqlite3
import tempfile

from behave import given, then, when

from cleveragents.infrastructure.sandbox.factory import SandboxFactory
from cleveragents.infrastructure.sandbox.protocol import SandboxStatus
from cleveragents.infrastructure.sandbox.transaction_sandbox import (
    TransactionSandbox,
)
from cleveragents.resource.handlers.database import (
    DATABASE_TYPE_DEFS,
    DatabaseResourceHandler,
    resolve_connection_args,
    validate_connection,
)
from cleveragents.resource.handlers.protocol import ResourceHandler

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _find_type_def(name):
    """Find a database type definition by name."""
    for td in DATABASE_TYPE_DEFS:
        if td["name"] == name:
            return td
    return None


# ---------------------------------------------------------------------------
# Resource Type Definitions
# ---------------------------------------------------------------------------


@given("the database type definitions are loaded")
def given_db_type_defs_loaded(context):
    """Load database type definitions into context."""
    context.db_type_defs = {td["name"]: td for td in DATABASE_TYPE_DEFS}


@then('the "{name}" type definition should exist')
def then_type_def_exists(context, name):
    assert name in context.db_type_defs, f"Type '{name}' not found in definitions"


@then('the "{name}" type should have sandbox strategy "{strategy}"')
def then_type_strategy(context, name, strategy):
    assert context.db_type_defs[name]["sandbox_strategy"] == strategy


@then('the "{name}" type should have a "{arg_name}" cli argument')
def then_type_has_arg(context, name, arg_name):
    cli_args = context.db_type_defs[name]["cli_args"]
    arg_names = [a["name"] for a in cli_args]
    assert arg_name in arg_names, f"Argument '{arg_name}' not found in {arg_names}"


# ---------------------------------------------------------------------------
# Connection Argument Resolution
# ---------------------------------------------------------------------------


@given(
    'postgres connection properties with host "{host}" port {port:d} dbname "{dbname}"'
)
def given_pg_props(context, host, port, dbname):
    context.conn_props = {"host": host, "port": port, "dbname": dbname}


@given('postgres properties with connection string "{conn_str}"')
def given_pg_conn_str(context, conn_str):
    context.conn_props = {"connection-string": conn_str}


@given("postgres properties without explicit credentials")
def given_pg_no_creds(context):
    context.conn_props = {"host": "localhost", "port": 5432}


@given('sqlite properties with path "{path}"')
def given_sqlite_props(context, path):
    context.conn_props = {"path": path}


@given("duckdb properties with no path")
def given_duckdb_no_path(context):
    context.conn_props = {}


@when('I resolve connection args for "{resource_type}"')
def when_resolve_args(context, resource_type):
    context.resolved_args = resolve_connection_args(context.conn_props, resource_type)


@then('the resolved args should have host "{host}"')
def then_resolved_host(context, host):
    assert context.resolved_args["host"] == host


@then("the resolved args should have port {port:d}")
def then_resolved_port(context, port):
    assert context.resolved_args["port"] == port


@then('the resolved args should have dbname "{dbname}"')
def then_resolved_dbname(context, dbname):
    assert context.resolved_args["dbname"] == dbname


@then('the resolved args should have connection string "{conn_str}"')
def then_resolved_conn_str(context, conn_str):
    assert context.resolved_args["connection-string"] == conn_str


@then('the resolved args should have user "{user}"')
def then_resolved_user(context, user):
    args = context.resolved_args
    assert args["user"] == user, f"Expected user '{user}', got '{args['user']}'"
    _cleanup_env(context)


@then('the resolved args should have password "{password}"')
def then_resolved_password(context, password):
    args = context.resolved_args
    assert args["password"] == password
    _cleanup_env(context)


@then('the resolved args should have path "{path}"')
def then_resolved_path(context, path):
    assert context.resolved_args["path"] == path


def _cleanup_env(context):
    """Restore saved environment variables."""
    saved = getattr(context, "_saved_env", {})
    for var, old_value in saved.items():
        if old_value is None:
            os.environ.pop(var, None)
        else:
            os.environ[var] = old_value
    context._saved_env = {}


# ---------------------------------------------------------------------------
# Credential Masking
# ---------------------------------------------------------------------------


@given('postgres connection properties with user "{user}" and password "{password}"')
def given_pg_with_creds(context, user, password):
    context.conn_props = {
        "host": "localhost",
        "user": user,
        "password": password,
    }


@when('I validate the "{resource_type}" connection')
def when_validate_connection(context, resource_type):
    resolved = resolve_connection_args(context.conn_props, resource_type)
    success, message = validate_connection(resource_type, resolved)
    context.validation_success = success
    context.validation_message = message


@then('the validation message should not contain "{secret}"')
def then_message_no_secret(context, secret):
    msg = context.validation_message
    assert secret not in msg, f"Secret '{secret}' found in validation message: {msg}"


@then("the validation result should be successful")
def then_validation_success(context):
    assert context.validation_success is True


# ---------------------------------------------------------------------------
# Read-Only Enforcement
# ---------------------------------------------------------------------------


@given("a temporary SQLite database with a test table")
def given_temp_sqlite_db(context):
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    conn = sqlite3.connect(path)
    conn.execute("CREATE TABLE test_data (id INTEGER PRIMARY KEY, value TEXT)")
    conn.commit()
    conn.close()
    context.db_path = path


@given("a transaction sandbox in read-only mode")
def given_ro_sandbox(context):
    sandbox = TransactionSandbox(
        resource_id="res-ro-1",
        original_path=context.db_path,
        read_only=True,
    )
    sandbox.create(plan_id="PLAN-RO")
    context.tx_sandbox = sandbox


@when("I attempt to insert data via the sandbox")
def when_insert_readonly(context):
    sandbox = context.tx_sandbox
    try:
        sandbox.execute("INSERT INTO test_data (id, value) VALUES (1, 'forbidden')")
        context.write_rejected = False
    except Exception:
        context.write_rejected = True
    finally:
        sandbox.cleanup()


@then("the sandbox should reject the write operation")
def then_write_rejected(context):
    assert context.write_rejected is True


# ---------------------------------------------------------------------------
# Transaction Sandbox Lifecycle
# ---------------------------------------------------------------------------


@given("a transaction sandbox for the database")
def given_tx_sandbox(context):
    sandbox = TransactionSandbox(
        resource_id="res-tx-1",
        original_path=context.db_path,
    )
    context.tx_sandbox = sandbox


@when('I create the sandbox with plan "{plan_id}"')
def when_create_sandbox(context, plan_id):
    context.tx_sandbox.create(plan_id=plan_id)


@when("I insert a row via the sandbox")
def when_insert_row(context):
    context.tx_sandbox.execute(
        "INSERT INTO test_data (id, value) VALUES (99, 'test-value')"
    )


@when("I commit the sandbox")
def when_commit_sandbox(context):
    context.commit_result = context.tx_sandbox.commit("test commit")
    context.tx_sandbox.cleanup()


@when("I rollback the sandbox")
def when_rollback_sandbox(context):
    context.tx_sandbox.rollback()
    context.tx_sandbox.cleanup()


@when("I cleanup the sandbox")
def when_cleanup_sandbox(context):
    context.tx_sandbox.cleanup()


@then("the row should be persisted in the database")
def then_row_persisted(context):
    conn = sqlite3.connect(context.db_path)
    rows = conn.execute("SELECT value FROM test_data WHERE id = 99").fetchall()
    conn.close()
    assert len(rows) == 1, f"Expected 1 row, got {len(rows)}"
    assert rows[0][0] == "test-value"


@then("the row should not be persisted in the database")
def then_row_not_persisted(context):
    conn = sqlite3.connect(context.db_path)
    rows = conn.execute("SELECT value FROM test_data WHERE id = 99").fetchall()
    conn.close()
    assert len(rows) == 0, f"Expected 0 rows, got {len(rows)}"


@then('the sandbox status should be "{status}"')
def then_sandbox_status(context, status):
    assert context.tx_sandbox.status == SandboxStatus(status)


# ---------------------------------------------------------------------------
# SQLite Validation
# ---------------------------------------------------------------------------


@given("a temporary SQLite database")
def given_temp_sqlite(context):
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    context.sqlite_path = path


@when("I validate the sqlite connection")
def when_validate_sqlite(context):
    success, msg = validate_connection("sqlite", {"path": context.sqlite_path})
    context.sqlite_valid_ok = success
    context.sqlite_valid_msg = msg


@when("I validate an in-memory sqlite connection")
def when_validate_sqlite_memory(context):
    success, msg = validate_connection("sqlite", {"path": ":memory:"})
    context.sqlite_valid_ok = success
    context.sqlite_valid_msg = msg


@then("the sqlite validation should succeed")
def then_sqlite_valid(context):
    assert context.sqlite_valid_ok is True


# ---------------------------------------------------------------------------
# DatabaseResourceHandler Protocol
# ---------------------------------------------------------------------------


@given("a DatabaseResourceHandler instance")
def given_db_handler(context):
    context.db_handler = DatabaseResourceHandler()


@then("the handler should satisfy the ResourceHandler protocol")
def then_handler_protocol(context):
    assert isinstance(context.db_handler, ResourceHandler)


@then('the handler default strategy should be "{strategy}"')
def then_handler_strategy(context, strategy):
    assert context.db_handler._default_strategy.value == strategy


# ---------------------------------------------------------------------------
# SandboxFactory transaction_rollback routing
# ---------------------------------------------------------------------------


@when(
    'a sandbox is requested for resource "{res_id}" '
    'at "{path}" using transaction_rollback'
)
def when_factory_tx(context, res_id, path):
    context.factory_sandbox = context.sandbox_factory.create_sandbox(
        resource_id=res_id,
        original_path=path,
        sandbox_strategy="transaction_rollback",
    )


@then("the created sandbox should be a TransactionSandbox")
def then_sandbox_is_tx(context):
    assert isinstance(context.factory_sandbox, TransactionSandbox)


@then('the factory should report "{strategy}" as supported')
def then_factory_supported(context, strategy):
    assert SandboxFactory.is_supported(strategy) is True


# ---------------------------------------------------------------------------
# TransactionSandbox edge cases
# ---------------------------------------------------------------------------


@when("I create a TransactionSandbox with empty resource_id")
def when_tx_empty_res_id(context):
    try:
        TransactionSandbox(resource_id="", original_path="/tmp/x")
        context.edge_error = None
    except ValueError as exc:
        context.edge_error = exc


@when("I create a TransactionSandbox with empty original_path")
def when_tx_empty_path(context):
    try:
        TransactionSandbox(resource_id="res-1", original_path="")
        context.edge_error = None
    except ValueError as exc:
        context.edge_error = exc


@then('the edge error should be a ValueError mentioning "{fragment}"')
def then_value_error_with_msg(context, fragment):
    assert context.edge_error is not None, "Expected a ValueError"
    assert isinstance(context.edge_error, ValueError)
    assert fragment in str(context.edge_error)


@when("I try to create the sandbox with an empty plan_id")
def when_empty_plan_id(context):
    try:
        context.tx_sandbox.create(plan_id="")
        context.edge_error = None
    except ValueError as exc:
        context.edge_error = exc


@then("get_path should return the database path")
def then_get_path_returns_db(context):
    sandbox = context.tx_sandbox
    result = sandbox.get_path("table")
    assert result == context.db_path


@then("get_path with traversal should raise ValueError")
def then_get_path_traversal(context):
    sandbox = context.tx_sandbox
    raised = False
    try:
        sandbox.get_path("../../etc/passwd")
    except ValueError:
        raised = True
    assert raised, "Expected ValueError for path traversal"


@then("get_path before create should raise SandboxStateError")
def then_get_path_wrong_state(context):
    from cleveragents.infrastructure.sandbox.protocol import SandboxStateError

    sandbox = context.tx_sandbox
    raised = False
    try:
        sandbox.get_path("test")
    except SandboxStateError:
        raised = True
    assert raised, "Expected SandboxStateError"


@then("commit before create should raise SandboxStateError")
def then_commit_wrong_state(context):
    from cleveragents.infrastructure.sandbox.protocol import SandboxStateError

    sandbox = context.tx_sandbox
    raised = False
    try:
        sandbox.commit("test")
    except SandboxStateError:
        raised = True
    assert raised, "Expected SandboxStateError"


@then("rollback before create should raise SandboxStateError")
def then_rollback_wrong_state(context):
    from cleveragents.infrastructure.sandbox.protocol import SandboxStateError

    sandbox = context.tx_sandbox
    raised = False
    try:
        sandbox.rollback()
    except SandboxStateError:
        raised = True
    assert raised, "Expected SandboxStateError"


@when("I cleanup the sandbox again")
def when_cleanup_again(context):
    context.tx_sandbox.cleanup()


@then("execute before create should raise SandboxStateError")
def then_execute_wrong_state(context):
    from cleveragents.infrastructure.sandbox.protocol import SandboxStateError

    sandbox = context.tx_sandbox
    raised = False
    try:
        sandbox.execute("SELECT 1")
    except SandboxStateError:
        raised = True
    assert raised, "Expected SandboxStateError"


@when('I validate a duckdb connection with path "{path}"')
def when_validate_duckdb(context, path):
    success, msg = validate_connection("duckdb", {"path": path})
    context.duckdb_valid_ok = success
    context.duckdb_valid_msg = msg


@then("the duckdb validation should succeed with a stub message")
def then_duckdb_stub(context):
    assert context.duckdb_valid_ok is True
    # Either driver works or stub message about not installed
    assert (
        "DuckDB" in context.duckdb_valid_msg
        or "duckdb" in context.duckdb_valid_msg.lower()
    )


@given('mysql connection properties with host "{host}" and port {port:d}')
def given_mysql_props(context, host, port):
    context.conn_props = {"host": host, "port": port}


@when("I validate an unknown database type")
def when_validate_unknown_db(context):
    success, msg = validate_connection("unknown_db", {})
    context.unknown_valid_ok = success
    context.unknown_valid_msg = msg


@then("the unknown validation should fail")
def then_unknown_fail(context):
    assert context.unknown_valid_ok is False
    assert "Unknown" in context.unknown_valid_msg


@given('mysql properties with connection string "{conn_str}"')
def given_mysql_conn_str(context, conn_str):
    context.conn_props = {"connection-string": conn_str}