devplacepy/devplacepy_services/base/db_codec.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

from datetime import date, datetime
from pathlib import Path
from typing import Any

WRITE_PREFIXES = (
    "set_",
    "add_",
    "delete_",
    "create_",
    "update_",
    "upsert_",
    "record_",
    "migrate_",
    "backfill_",
    "mark_",
    "invalidate_",
    "ensure_",
    "restore",
    "purge",
    "revoke_",
    "recompute_",
    "soft_delete",
    "bump_",
    "init_",
    "clear_settings",
)

WRITE_EXACT = frozenset(
    {
        "soft_delete_in",
        "restore_event",
        "purge_event",
        "delete_engagement",
        "soft_delete_engagement",
        "delete_fork_relations",
        "soft_delete_fork_relations",
        "delete_attachments",
        "delete_attachment_record",
        "_delete_attachment_file",
        "_index",
        "_drop_index",
        "_uid_index",
        "_ensure_cache_state",
        "_refresh_query_planner_stats",
        "_backfill_gamification",
    }
)

REMOTE_TABLE_MARKER = "__remote_table__"


def is_write(name: str) -> bool:
    if name in WRITE_EXACT:
        return True
    return any(name.startswith(prefix) for prefix in WRITE_PREFIXES)


SQL_WRITE_KEYWORDS = frozenset(
    {"INSERT", "UPDATE", "DELETE", "REPLACE", "CREATE", "ALTER", "DROP"}
)


def is_write_sql(sql: str) -> bool:
    first_word = sql.strip().split(None, 1)[0].upper() if sql.strip() else ""
    return first_word in SQL_WRITE_KEYWORDS


def encode_value(value: Any) -> Any:
    if value is None or isinstance(value, (bool, int, float, str)):
        return value
    if type(value).__name__ == "RemoteTable":
        return {REMOTE_TABLE_MARKER: value._name}
    if isinstance(value, (datetime, date)):
        return value.isoformat()
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, frozenset):
        return [encode_value(item) for item in value]
    if isinstance(value, set):
        return [encode_value(item) for item in value]
    if isinstance(value, tuple):
        return [encode_value(item) for item in value]
    if isinstance(value, list):
        return [encode_value(item) for item in value]
    if isinstance(value, dict):
        return {str(key): encode_value(item) for key, item in value.items()}
    if hasattr(value, "items") and callable(value.items):
        try:
            return {str(key): encode_value(item) for key, item in value.items()}
        except TypeError:
            pass
    return str(value)


def encode_args(args: tuple | list, kwargs: dict) -> tuple[list, dict]:
    return [encode_value(item) for item in args], {
        str(key): encode_value(value) for key, value in kwargs.items()
    }


def decode_value(value: Any) -> Any:
    if isinstance(value, list):
        return [decode_value(item) for item in value]
    if isinstance(value, dict):
        return {key: decode_value(item) for key, item in value.items()}
    return value


def decode_arg(value: Any) -> Any:
    if isinstance(value, list):
        return [decode_arg(item) for item in value]
    if isinstance(value, dict):
        if set(value) == {REMOTE_TABLE_MARKER}:
            from devplacepy.db_client import get_table

            return get_table(value[REMOTE_TABLE_MARKER])
        return {key: decode_arg(item) for key, item in value.items()}
    return value


def encode_result(value: Any) -> Any:
    return encode_value(value)