minidb-sql 1.3.0__py3-none-any.whl

minidb_sql/table.py

@@ -0,0 +1,274 @@
+"""Table storage and row management for MiniDB."""
+
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass, field
+from typing import Any
+
+from .column import Schema
+from .errors import ColumnNotFoundError, DuplicateKeyError
+from .index import IndexManager
+from .types import Row
+
+
+@dataclass
+class Table:
+    """
+    Represents a database table with rows and a schema.
+
+    Stores rows as dictionaries and maintains indexes for fast lookups.
+    """
+
+    name: str
+    schema: Schema
+    _rows: list[Row] = field(default_factory=list)
+    _row_ids: list[int] = field(default_factory=list)
+    _next_row_id: int = field(default=0)
+    _index_manager: IndexManager = field(default_factory=IndexManager)
+
+    def __post_init__(self):
+        """Initialize indexes after table creation."""
+        # Create index on primary key if exists
+        if self.schema.primary_key:
+            self._index_manager.create_index(self.schema.primary_key)
+
+    @property
+    def row_count(self) -> int:
+        """Get the number of rows in the table."""
+        return len(self._rows)
+
+    @property
+    def columns(self) -> list[str]:
+        """Get list of column names."""
+        return self.schema.column_names
+
+    @property
+    def primary_key(self) -> str | None:
+        """Get the primary key column name."""
+        return self.schema.primary_key
+
+    def insert(self, row: dict) -> int:
+        """
+        Insert a new row into the table.
+
+        Args:
+            row: Row data as a dictionary
+
+        Returns:
+            The row ID of the inserted row
+
+        Raises:
+            DuplicateKeyError: If primary key already exists
+        """
+        # Validate and normalize the row
+        validated_row = self.schema.validate_row(row, allow_missing=True)
+
+        # Check for duplicate primary key
+        pk_column = self.schema.primary_key
+        if pk_column:
+            pk_value = validated_row[pk_column]
+            if pk_value is not None:
+                index = self._index_manager.get_index(pk_column)
+                if index and index.has_value(pk_value):
+                    raise DuplicateKeyError(pk_value)
+
+        # Assign row ID
+        row_id = self._next_row_id
+        self._next_row_id += 1
+
+        # Store the row
+        self._rows.append(validated_row)
+        self._row_ids.append(row_id)
+
+        # Update indexes
+        self._index_manager.insert_row(validated_row, row_id)
+
+        return row_id
+
+    def get_row_by_id(self, row_id: int) -> Row | None:
+        """Get a row by its ID."""
+        try:
+            idx = self._row_ids.index(row_id)
+            return self._rows[idx].copy()
+        except ValueError:
+            return None
+
+    def rows_for_ids(self, row_ids: Iterable[int]) -> list[tuple[int, Row]]:
+        """Fetch many rows by id in one pass.
+
+        get_row_by_id walks every stored id. A primary-key range that
+        matches thousands of rows was slower than reading the table once.
+        """
+        position = {row_id: index for index, row_id in enumerate(self._row_ids)}
+        found: list[tuple[int, Row]] = []
+        for row_id in row_ids:
+            index = position.get(row_id)
+            if index is not None:
+                found.append((row_id, self._rows[index].copy()))
+        return found
+
+    def get_row_id_by_pk(self, pk_value: Any) -> int | None:
+        """Get a row ID by primary key value."""
+        pk_column = self.schema.primary_key
+        if not pk_column:
+            return None
+
+        index = self._index_manager.get_index(pk_column)
+        if not index:
+            return None
+
+        row_ids = index.lookup(pk_value)
+        return next(iter(row_ids)) if row_ids else None
+
+    def get_all_rows(self) -> list[Row]:
+        """Get all rows in the table."""
+        return [row.copy() for row in self._rows]
+
+    def get_rows_with_ids(self) -> list[tuple[int, Row]]:
+        """Get all rows with their IDs."""
+        return [(row_id, row.copy()) for row_id, row in zip(self._row_ids, self._rows, strict=False)]
+
+    def update_rows(self, row_ids: set[int], updates: dict) -> int:
+        """
+        Update rows with the given IDs.
+
+        Args:
+            row_ids: Set of row IDs to update
+            updates: Dictionary of column updates
+
+        Returns:
+            Number of rows updated
+        """
+        if not row_ids or not updates:
+            return 0
+
+        # Plan every row before writing. A later duplicate must not keep
+        # the first row's new primary key.
+        planned: list[tuple[int, int, dict, dict]] = []
+        pk_column = self.schema.primary_key
+        new_keys: dict[Any, int] = {}
+        for i, row_id in enumerate(self._row_ids):
+            if row_id not in row_ids:
+                continue
+            old_row = self._rows[i]
+            new_row = old_row.copy()
+            for col, val in updates.items():
+                col_def = self.schema.get_column(col)
+                if col_def is None:
+                    raise ColumnNotFoundError(col, self.name)
+                new_row[col] = col_def.validate(val)
+            validated = self.schema.validate_row(new_row, allow_missing=True)
+            if pk_column and pk_column in updates:
+                new_pk = validated[pk_column]
+                if new_pk is not None:
+                    if new_pk in new_keys:
+                        raise DuplicateKeyError(new_pk)
+                    index = self.get_index(pk_column)
+                    if index and index.has_value(new_pk):
+                        existing_ids = index.lookup(new_pk)
+                        if any(existing_id not in row_ids for existing_id in existing_ids):
+                            raise DuplicateKeyError(new_pk)
+                    new_keys[new_pk] = row_id
+            planned.append((i, row_id, old_row, validated))
+
+        for i, row_id, old_row, validated in planned:
+            self._index_manager.update_row(old_row, validated, row_id)
+            self._rows[i] = validated
+
+        return len(planned)
+
+    def delete_rows(self, row_ids: set[int]) -> int:
+        """
+        Delete rows with the given IDs.
+
+        Args:
+            row_ids: Set of row IDs to delete
+
+        Returns:
+            Number of rows deleted
+        """
+        if not row_ids:
+            return 0
+
+        # Find indices to delete (in reverse order for safe removal)
+        indices_to_delete = []
+        for i, row_id in enumerate(self._row_ids):
+            if row_id in row_ids:
+                indices_to_delete.append(i)
+
+        # Delete in reverse order
+        count = 0
+        for i in reversed(indices_to_delete):
+            row = self._rows[i]
+            row_id = self._row_ids[i]
+
+            # Update indexes
+            self._index_manager.delete_row(row, row_id)
+
+            # Remove row
+            self._rows.pop(i)
+            self._row_ids.pop(i)
+            count += 1
+
+        return count
+
+    def scan(self) -> Iterator[tuple[int, Row]]:
+        """Iterate over all rows with their IDs."""
+        for row_id, row in zip(self._row_ids, self._rows, strict=False):
+            yield row_id, row.copy()
+
+    def get_index(self, column_name: str):
+        """Get an index for a column."""
+        return self._index_manager.get_index(column_name)
+
+    def has_index(self, column_name: str) -> bool:
+        """Check if an index exists for a column."""
+        return self._index_manager.has_index(column_name)
+
+    def clear(self) -> None:
+        """Remove all rows from the table."""
+        self._rows.clear()
+        self._row_ids.clear()
+        self._next_row_id = 0
+        self._index_manager.clear()
+
+    def to_dict(self) -> dict:
+        """Serialize table to a dictionary."""
+        return {
+            'name': self.name,
+            'schema': self.schema.to_dict(),
+            'rows': self._rows,
+            'row_ids': list(self._row_ids),
+            'next_row_id': self._next_row_id,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> 'Table':
+        """Deserialize table from a dictionary."""
+        schema = Schema.from_dict(data['schema'])
+        table = cls(name=data['name'], schema=schema)
+        # Same checks as insert. A hand-edited file must not keep two primary
+        # keys or a value the column cannot store.
+        table._rows = [schema.validate_row(row, allow_missing=True) for row in data['rows']]
+        table._next_row_id = data.get('next_row_id', len(table._rows))
+
+        # Files written before row_ids existed keep 0..n-1.
+        saved_ids = data.get('row_ids')
+        if isinstance(saved_ids, list) and len(saved_ids) == len(table._rows):
+            table._row_ids = list(saved_ids)
+        else:
+            table._row_ids = list(range(len(table._rows)))
+
+        pk = table.schema.primary_key
+        if pk:
+            index = table._index_manager.get_index(pk)
+            if index is None:
+                index = table._index_manager.create_index(pk)
+            for row, row_id in zip(table._rows, table._row_ids, strict=True):
+                value = row.get(pk)
+                if value is None:
+                    continue
+                if index.has_value(value):
+                    raise DuplicateKeyError(value)
+                index.insert(value, row_id)
+
+        return table

minidb_sql/types.py

@@ -0,0 +1,189 @@
+"""Type definitions for MiniDB."""
+
+from enum import Enum, auto
+from typing import Any
+
+
+class ColumnType(Enum):
+    """Supported column types in MiniDB."""
+
+    INTEGER = auto()
+    STRING = auto()
+    FLOAT = auto()
+    BOOLEAN = auto()
+
+    @classmethod
+    def from_string(cls, type_str: str) -> 'ColumnType':
+        """Convert a string representation to ColumnType."""
+        type_map = {
+            'INTEGER': cls.INTEGER,
+            'INT': cls.INTEGER,
+            'STRING': cls.STRING,
+            'STR': cls.STRING,
+            'TEXT': cls.STRING,
+            'VARCHAR': cls.STRING,
+            'FLOAT': cls.FLOAT,
+            'DOUBLE': cls.FLOAT,
+            'REAL': cls.FLOAT,
+            'BOOLEAN': cls.BOOLEAN,
+            'BOOL': cls.BOOLEAN,
+        }
+        upper_str = type_str.upper()
+        if upper_str not in type_map:
+            raise ValueError(f'Unknown column type: {type_str}')
+        return type_map[upper_str]
+
+    def to_string(self) -> str:
+        """Convert ColumnType to string representation."""
+        return self.name
+
+    def validate(self, value: Any) -> bool:
+        """Check if a value is valid for this column type."""
+        if value is None:
+            return True  # NULL values handled by nullable constraint
+
+        type_validators = {
+            ColumnType.INTEGER: lambda v: isinstance(v, int) and not isinstance(v, bool),
+            ColumnType.STRING: lambda v: isinstance(v, str),
+            ColumnType.FLOAT: lambda v: isinstance(v, (int, float)) and not isinstance(v, bool),
+            ColumnType.BOOLEAN: lambda v: isinstance(v, bool),
+        }
+        return type_validators[self](value)
+
+    def cast(self, value: Any) -> Any:
+        """Attempt to cast a value to this type."""
+        if value is None:
+            return None
+
+        try:
+            if self == ColumnType.INTEGER:
+                if isinstance(value, bool):
+                    raise TypeError(f'Cannot cast {value!r} to INTEGER')
+                if isinstance(value, float):
+                    if not value.is_integer():
+                        raise TypeError(f'Cannot cast {value!r} to INTEGER')
+                    return int(value)
+                if isinstance(value, int):
+                    return value
+                return int(value)
+            elif self == ColumnType.STRING:
+                return str(value)
+            elif self == ColumnType.FLOAT:
+                if isinstance(value, bool):
+                    raise TypeError(f'Cannot cast {value!r} to FLOAT')
+                return float(value)
+            elif self == ColumnType.BOOLEAN:
+                if isinstance(value, bool):
+                    return value
+                if isinstance(value, str):
+                    spelling = {
+                        'true': True,
+                        'false': False,
+                        '1': True,
+                        '0': False,
+                        'yes': True,
+                        'no': False,
+                    }
+                    if value.lower() not in spelling:
+                        raise TypeError(f'Cannot cast {value!r} to BOOLEAN')
+                    return spelling[value.lower()]
+                raise TypeError(f'Cannot cast {value!r} to BOOLEAN')
+        except (ValueError, TypeError) as e:
+            raise TypeError(f'Cannot cast {value!r} to {self.name}') from e
+
+        return value
+
+
+class TokenType(Enum):
+    """Token types for the SQL lexer."""
+
+    # Keywords
+    SELECT = auto()
+    FROM = auto()
+    WHERE = auto()
+    INSERT = auto()
+    INTO = auto()
+    VALUES = auto()
+    UPDATE = auto()
+    SET = auto()
+    DELETE = auto()
+    CREATE = auto()
+    TABLE = auto()
+    DROP = auto()
+    ORDER = auto()
+    BY = auto()
+    GROUP = auto()
+    JOIN = auto()
+    INNER = auto()
+    LEFT = auto()
+    RIGHT = auto()
+    ON = auto()
+    AND = auto()
+    OR = auto()
+    NOT = auto()
+    IN = auto()
+    LIKE = auto()
+    ASC = auto()
+    DESC = auto()
+    PRIMARY = auto()
+    KEY = auto()
+    NULL = auto()
+    NOT_NULL = auto()
+    IS = auto()
+    LIMIT = auto()
+    AS = auto()
+    DISTINCT = auto()
+    HAVING = auto()
+    CASE = auto()
+    WHEN = auto()
+    THEN = auto()
+    ELSE = auto()
+    END = auto()
+    BEGIN = auto()
+    COMMIT = auto()
+    ROLLBACK = auto()
+
+    # Literals
+    IDENTIFIER = auto()
+    STRING_LITERAL = auto()
+    INTEGER_LITERAL = auto()
+    FLOAT_LITERAL = auto()
+    BOOLEAN_LITERAL = auto()
+
+    # Operators
+    EQUALS = auto()
+    NOT_EQUALS = auto()
+    GREATER = auto()
+    GREATER_EQUAL = auto()
+    LESS = auto()
+    LESS_EQUAL = auto()
+
+    # Punctuation
+    COMMA = auto()
+    DOT = auto()
+    LPAREN = auto()
+    RPAREN = auto()
+    SEMICOLON = auto()
+    STAR = auto()
+    PLACEHOLDER = auto()
+
+    # Special
+    EOF = auto()
+
+
+class QueryType(Enum):
+    """Types of SQL queries."""
+
+    SELECT = auto()
+    INSERT = auto()
+    UPDATE = auto()
+    DELETE = auto()
+    CREATE_TABLE = auto()
+    DROP_TABLE = auto()
+
+
+# Type alias for row data
+Row = dict[str, Any]
+
+# Type alias for query result
+QueryResult = list[Row] | int | None