user_tables.py

"""
Hosts the table tiers, user tables should be derived from.
"""

import re

from .autopopulate import AutoPopulate
from .errors import DataJointError
from .table import Table
from .utils import from_camel_case

_base_regexp = r"[a-z][a-z0-9]*(_[a-z][a-z0-9]*)*"

# attributes that trigger instantiation of user classes


supported_class_attrs = {
    "key_source",
    "describe",
    "alter",
    "heading",
    "populate",
    "progress",
    "primary_key",
    "proj",
    "aggr",
    "join",
    "extend",
    "to_dicts",
    "to_pandas",
    "to_polars",
    "to_arrow",
    "to_arrays",
    "keys",
    "fetch",
    "fetch1",
    "head",
    "tail",
    "descendants",
    "ancestors",
    "parts",
    "parents",
    "children",
    "insert",
    "insert1",
    "insert_dataframe",
    "update1",
    "validate",
    "drop",
    "drop_quick",
    "delete",
    "delete_quick",
    "staged_insert1",
}


class TableMeta(type):
    """
    TableMeta subclasses allow applying some instance methods and properties directly
    at class level. For example, this allows Table.to_dicts() instead of Table().to_dicts().
    """

    def __getattribute__(cls, name):
        # trigger instantiation for supported class attrs
        return cls().__getattribute__(name) if name in supported_class_attrs else super().__getattribute__(name)

    def __and__(cls, arg):
        return cls() & arg

    def __xor__(cls, arg):
        return cls() ^ arg

    def __sub__(cls, arg):
        return cls() - arg

    def __neg__(cls):
        return -cls()

    def __mul__(cls, arg):
        return cls() * arg

    def __matmul__(cls, arg):
        return cls() @ arg

    def __add__(cls, arg):
        return cls() + arg

    def __iter__(cls):
        return iter(cls())

    # Class properties - defined on metaclass to work at class level
    @property
    def connection(cls):
        """The database connection for this table."""
        return cls._connection

    @property
    def table_name(cls):
        """The table name formatted for MySQL."""
        if cls._prefix is None:
            raise AttributeError("Class prefix is not defined!")
        return cls._prefix + from_camel_case(cls.__name__)

    @property
    def full_table_name(cls):
        """The fully qualified table name (quoted per backend)."""
        if cls.database is None:
            return None
        return cls._connection.adapter.make_full_table_name(cls.database, cls.table_name)


class UserTable(Table, metaclass=TableMeta):
    """
    A subclass of UserTable is a dedicated class interfacing a base table.
    UserTable is initialized by the decorator generated by schema().
    """

    # set by @schema
    _connection = None
    _heading = None
    _support = None

    # set by subclass
    tier_regexp = None
    _prefix = None

    @property
    def definition(self):
        """
        :return: a string containing the table definition using the DataJoint DDL.
        """
        raise NotImplementedError('Subclasses of Table must implement the property "definition"')


class Manual(UserTable):
    """
    Inherit from this class if the table's values are entered manually.
    """

    _prefix = r""
    tier_regexp = r"(?P<manual>" + _prefix + _base_regexp + ")"


class Lookup(UserTable):
    """
    Inherit from this class if the table's values are for lookup. This is
    currently equivalent to defining the table as Manual and serves semantic
    purposes only.
    """

    _prefix = "#"
    tier_regexp = r"(?P<lookup>" + _prefix + _base_regexp.replace("TIER", "lookup") + ")"


class Imported(UserTable, AutoPopulate):
    """
    Inherit from this class if the table's values are imported from external data sources.
    The inherited class must at least provide the function `_make_tuples`.
    """

    _prefix = "_"
    tier_regexp = r"(?P<imported>" + _prefix + _base_regexp + ")"


class Computed(UserTable, AutoPopulate):
    """
    Inherit from this class if the table's values are computed from other tables in the schema.
    The inherited class must at least provide the function `_make_tuples`.
    """

    _prefix = "__"
    tier_regexp = r"(?P<computed>" + _prefix + _base_regexp + ")"


class PartMeta(TableMeta):
    """Metaclass for Part tables with overridden class properties."""

    @property
    def table_name(cls):
        """The table name for a Part is derived from its master table."""
        return None if cls.master is None else cls.master.table_name + "__" + from_camel_case(cls.__name__)

    @property
    def full_table_name(cls):
        """The fully qualified table name (quoted per backend)."""
        if cls.database is None or cls.table_name is None:
            return None
        return cls._connection.adapter.make_full_table_name(cls.database, cls.table_name)

    @property
    def master(cls):
        """The master table for this Part table."""
        return cls._master


class Part(UserTable, metaclass=PartMeta):
    """
    Inherit from this class if the table's values are details of an entry in another table
    and if this table is populated by the other table. For example, the entries inheriting from
    dj.Part could be single entries of a matrix, while the parent table refers to the entire matrix.
    Part tables are implemented as classes inside classes.
    """

    _connection = None
    _master = None

    tier_regexp = (
        r"(?P<master>"
        + "|".join([c.tier_regexp for c in (Manual, Lookup, Imported, Computed)])
        + r"){1,1}"
        + "__"
        + r"(?P<part>"
        + _base_regexp
        + ")"
    )

    def delete(self, part_integrity: str = "enforce", **kwargs):
        """
        Delete from a Part table.

        Args:
            part_integrity: Policy for master-part integrity. One of:
                - ``"enforce"`` (default): Error - delete from master instead.
                - ``"ignore"``: Allow direct deletion (breaks master-part integrity).
                - ``"cascade"``: Delete parts AND cascade up to delete master.
            **kwargs: Additional arguments passed to Table.delete()
                      (transaction, prompt)

        Raises:
            DataJointError: If part_integrity="enforce" (direct Part deletes prohibited)
        """
        if part_integrity == "enforce":
            raise DataJointError(
                "Cannot delete from a Part directly. Delete from master instead, "
                "or use part_integrity='ignore' to break integrity, "
                "or part_integrity='cascade' to also delete master."
            )
        super().delete(part_integrity=part_integrity, **kwargs)

    def drop(self, part_integrity: str = "enforce"):
        """
        Drop a Part table.

        Args:
            part_integrity: Policy for master-part integrity. One of:
                - ``"enforce"`` (default): Error - drop master instead.
                - ``"ignore"``: Allow direct drop (breaks master-part structure).
                Note: ``"cascade"`` is not supported for drop (too destructive).

        Raises:
            DataJointError: If part_integrity="enforce" (direct Part drops prohibited)
        """
        if part_integrity == "ignore":
            super().drop()
        elif part_integrity == "enforce":
            raise DataJointError("Cannot drop a Part directly. Drop master instead, or use part_integrity='ignore' to force.")
        else:
            raise ValueError(f"part_integrity for drop must be 'enforce' or 'ignore', got {part_integrity!r}")

    def alter(self, prompt=True, context=None):
        # without context, use declaration context which maps master keyword to master table
        super().alter(prompt=prompt, context=context or self.declaration_context)


user_table_classes = (Manual, Lookup, Computed, Imported, Part)


class _AliasNode:
    """
    special class to indicate aliased foreign keys
    """

    pass


def _get_tier(table_name):
    """given the table name, return the user table class."""
    # Handle both MySQL backticks and PostgreSQL double quotes
    if table_name.startswith("`"):
        # MySQL format: `schema`.`table_name`
        extracted_name = table_name.split("`")[-2]
    elif table_name.startswith('"'):
        # PostgreSQL format: "schema"."table_name"
        extracted_name = table_name.split('"')[-2]
    else:
        return _AliasNode
    try:
        return next(tier for tier in user_table_classes if re.fullmatch(tier.tier_regexp, extracted_name))
    except StopIteration:
        return None