954 lines
41 KiB
Python
954 lines
41 KiB
Python
"""Data access layer for the bitable subsystem.
|
||
|
||
All PostgreSQL operations go through this repository. The service layer
|
||
(:mod:`agentkit.bitable.service`) calls these methods — routes never
|
||
access the repository directly.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import re
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
from sqlalchemy import delete, func, insert, select, text, update
|
||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||
|
||
from agentkit.bitable.db import (
|
||
BitableDB,
|
||
FieldModel,
|
||
FileModel,
|
||
RecordModel,
|
||
RecalcQueueModel,
|
||
TableModel,
|
||
ViewModel,
|
||
_uuid_str,
|
||
)
|
||
from agentkit.bitable.models import (
|
||
BitableFile,
|
||
Field,
|
||
FieldOwner,
|
||
FieldType,
|
||
Record,
|
||
RecalcStatus,
|
||
RecalcTask,
|
||
Table,
|
||
View,
|
||
ViewType,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# U3/DR-1: field_id 在 jsonb_set path 与 where 子句中走字符串插值（jsonb path 不支持绑定参数）。
|
||
# 防御性校验：field_id 必须匹配 UUID 格式，拒绝任意字符串插值到 SQL 结构。
|
||
# ponytail: ceiling — 仅覆盖 field_id；其余 SQL 结构（op / direction）已由 service 层白名单枚举校验。
|
||
# 升级路径：迁移到 ORM JSON 函数（sqlalchemy.dialects.postgresql.jsonb_*）彻底消除字符串插值。
|
||
_FIELD_ID_RE = re.compile(
|
||
r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$"
|
||
)
|
||
|
||
|
||
def _validate_field_id(field_id: str, *, context: str = "") -> str:
|
||
"""校验 field_id 是 UUID 格式（DR-1 防御性校验）。
|
||
|
||
field_id 在 jsonb_set path 与 SQL where 子句中走字符串插值（jsonb path 不支持绑定参数），
|
||
必须确保非系统 UUID 不进入 SQL 结构。校验失败抛 ValueError。
|
||
"""
|
||
if not isinstance(field_id, str) or not _FIELD_ID_RE.match(field_id):
|
||
ctx = f" ({context})" if context else ""
|
||
raise ValueError(f"Invalid field_id format{ctx}: {field_id!r} — expected UUID")
|
||
return field_id
|
||
|
||
|
||
class BitableRepository:
|
||
"""Async repository for bitable CRUD operations.
|
||
|
||
Usage::
|
||
|
||
repo = BitableRepository(db)
|
||
table = await repo.create_table(name="My Table")
|
||
"""
|
||
|
||
def __init__(self, db: BitableDB) -> None:
|
||
self._db = db
|
||
|
||
@property
|
||
def _session_factory(self):
|
||
return self._db.session_factory
|
||
|
||
# ── Files (R1) ─────────────────────────────────────────
|
||
|
||
async def create_file(
|
||
self,
|
||
name: str,
|
||
icon: str = "table",
|
||
description: str = "",
|
||
owner_user_id: str | None = None,
|
||
) -> BitableFile:
|
||
"""Create a new bitable file (top-level container)."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
insert(FileModel)
|
||
.values(
|
||
id=_uuid_str(),
|
||
name=name,
|
||
icon=icon,
|
||
description=description,
|
||
owner_user_id=owner_user_id,
|
||
)
|
||
.returning(FileModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalar_one()
|
||
await session.commit()
|
||
return BitableFile.model_validate(entity)
|
||
|
||
async def get_file(self, file_id: str) -> BitableFile | None:
|
||
"""Get a file by ID."""
|
||
async with self._session_factory() as session:
|
||
stmt = select(FileModel).where(FileModel.id == file_id)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
return BitableFile.model_validate(entity) if entity else None
|
||
|
||
async def list_files(self, owner_user_id: str | None = None) -> list[BitableFile]:
|
||
"""List all files (optionally filtered by owner)."""
|
||
async with self._session_factory() as session:
|
||
stmt = select(FileModel)
|
||
if owner_user_id:
|
||
stmt = stmt.where(FileModel.owner_user_id == owner_user_id)
|
||
stmt = stmt.order_by(FileModel.updated_at.desc())
|
||
result = await session.execute(stmt)
|
||
return [BitableFile.model_validate(e) for e in result.scalars().all()]
|
||
|
||
async def update_file(self, file_id: str, **kwargs: object) -> BitableFile | None:
|
||
"""Update a file's attributes."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
update(FileModel)
|
||
.where(FileModel.id == file_id)
|
||
.values(**kwargs)
|
||
.returning(FileModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
await session.commit()
|
||
return BitableFile.model_validate(entity) if entity else None
|
||
|
||
async def delete_file(self, file_id: str) -> bool:
|
||
"""Delete a file. Caller is responsible for cascading table deletes."""
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(delete(FileModel).where(FileModel.id == file_id))
|
||
await session.commit()
|
||
return result.rowcount > 0
|
||
|
||
async def list_tables_by_file(self, file_id: str) -> list[Table]:
|
||
"""List all tables belonging to a file."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
select(TableModel)
|
||
.where(TableModel.file_id == file_id)
|
||
.order_by(TableModel.created_at)
|
||
)
|
||
result = await session.execute(stmt)
|
||
return [Table.model_validate(e) for e in result.scalars().all()]
|
||
|
||
# ── Tables ──────────────────────────────────────────────
|
||
|
||
async def create_table(
|
||
self,
|
||
name: str,
|
||
description: str = "",
|
||
primary_key_field_id: str | None = None,
|
||
owner_user_id: str | None = None,
|
||
file_id: str | None = None,
|
||
) -> Table:
|
||
"""Create a new bitable table."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
insert(TableModel)
|
||
.values(
|
||
id=_uuid_str(),
|
||
file_id=file_id,
|
||
name=name,
|
||
description=description,
|
||
primary_key_field_id=primary_key_field_id,
|
||
owner_user_id=owner_user_id,
|
||
)
|
||
.returning(TableModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalar_one()
|
||
await session.commit()
|
||
return Table.model_validate(entity)
|
||
|
||
async def get_table(self, table_id: str) -> Table | None:
|
||
"""Get a table by ID."""
|
||
async with self._session_factory() as session:
|
||
stmt = select(TableModel).where(TableModel.id == table_id)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
return Table.model_validate(entity) if entity else None
|
||
|
||
async def list_tables(self, owner_user_id: str | None = None) -> list[Table]:
|
||
"""List all tables (optionally filtered by owner)."""
|
||
async with self._session_factory() as session:
|
||
stmt = select(TableModel)
|
||
if owner_user_id:
|
||
stmt = stmt.where(TableModel.owner_user_id == owner_user_id)
|
||
stmt = stmt.order_by(TableModel.created_at.desc())
|
||
result = await session.execute(stmt)
|
||
return [Table.model_validate(e) for e in result.scalars().all()]
|
||
|
||
async def update_table(self, table_id: str, **kwargs: object) -> Table | None:
|
||
"""Update a table's attributes."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
update(TableModel)
|
||
.where(TableModel.id == table_id)
|
||
.values(**kwargs)
|
||
.returning(TableModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
await session.commit()
|
||
return Table.model_validate(entity) if entity else None
|
||
|
||
async def create_pk_unique_index(self, table_id: str, pk_field_id: str) -> None:
|
||
"""Create a unique expression index on the PK field for this table.
|
||
|
||
Enforces that ``values->>pk_field_id`` is unique per table. Required
|
||
for correct upsert semantics (KTD8) — without it, duplicate PK values
|
||
can be silently inserted and ``find_record_by_pk`` does full scans.
|
||
"""
|
||
# ponytail: pk_field_id is a system UUID from TableModel, validated here
|
||
# before interpolation into a SQL identifier.
|
||
if not re.match(r"^[a-f0-9-]{36}$", pk_field_id):
|
||
raise ValueError(f"Invalid pk_field_id format: {pk_field_id}")
|
||
index_name = f"ix_bitable_records_pk_{table_id.replace('-', '_')}"
|
||
sql = text(
|
||
f"CREATE UNIQUE INDEX IF NOT EXISTS {index_name} "
|
||
f"ON bitable.bitable_records (table_id, (values->>'{pk_field_id}')) "
|
||
f"WHERE values ? :pk_path"
|
||
)
|
||
async with self._session_factory() as session:
|
||
await session.execute(sql, {"pk_path": pk_field_id})
|
||
await session.commit()
|
||
|
||
async def delete_table(self, table_id: str) -> bool:
|
||
"""Delete a table and all its fields, records, views, and recalc tasks."""
|
||
async with self._session_factory() as session:
|
||
await session.execute(delete(FieldModel).where(FieldModel.table_id == table_id))
|
||
await session.execute(delete(RecordModel).where(RecordModel.table_id == table_id))
|
||
await session.execute(delete(ViewModel).where(ViewModel.table_id == table_id))
|
||
await session.execute(
|
||
delete(RecalcQueueModel).where(RecalcQueueModel.table_id == table_id)
|
||
)
|
||
result = await session.execute(delete(TableModel).where(TableModel.id == table_id))
|
||
await session.commit()
|
||
return result.rowcount > 0
|
||
|
||
# ── Fields ──────────────────────────────────────────────
|
||
|
||
async def create_field(
|
||
self,
|
||
table_id: str,
|
||
name: str,
|
||
field_type: FieldType,
|
||
config: dict[str, object] | None = None,
|
||
owner: FieldOwner = FieldOwner.user,
|
||
) -> Field:
|
||
"""Create a new field in a table."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
insert(FieldModel)
|
||
.values(
|
||
id=_uuid_str(),
|
||
table_id=table_id,
|
||
name=name,
|
||
field_type=field_type.value,
|
||
config=config or {},
|
||
owner=owner.value,
|
||
)
|
||
.returning(FieldModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalar_one()
|
||
await session.commit()
|
||
return Field.model_validate(entity)
|
||
|
||
async def get_field(self, field_id: str) -> Field | None:
|
||
"""Get a field by ID."""
|
||
async with self._session_factory() as session:
|
||
stmt = select(FieldModel).where(FieldModel.id == field_id)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
return Field.model_validate(entity) if entity else None
|
||
|
||
async def list_fields(self, table_id: str) -> list[Field]:
|
||
"""List all fields in a table."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
select(FieldModel)
|
||
.where(FieldModel.table_id == table_id)
|
||
.order_by(FieldModel.created_at)
|
||
)
|
||
result = await session.execute(stmt)
|
||
return [Field.model_validate(e) for e in result.scalars().all()]
|
||
|
||
async def update_field(self, field_id: str, **kwargs: object) -> Field | None:
|
||
"""Update a field's attributes."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
update(FieldModel)
|
||
.where(FieldModel.id == field_id)
|
||
.values(**kwargs)
|
||
.returning(FieldModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
await session.commit()
|
||
return Field.model_validate(entity) if entity else None
|
||
|
||
async def delete_field(self, field_id: str) -> bool:
|
||
"""Delete a field."""
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(delete(FieldModel).where(FieldModel.id == field_id))
|
||
await session.commit()
|
||
return result.rowcount > 0
|
||
|
||
# ── Records ─────────────────────────────────────────────
|
||
|
||
async def create_record(self, table_id: str, values: dict[str, object] | None = None) -> Record:
|
||
"""Create a new record."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
insert(RecordModel)
|
||
.values(
|
||
id=_uuid_str(),
|
||
table_id=table_id,
|
||
values=values or {},
|
||
)
|
||
.returning(RecordModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalar_one()
|
||
await session.commit()
|
||
return Record.model_validate(entity)
|
||
|
||
async def create_records_batch(
|
||
self, table_id: str, records_values: list[dict[str, object]]
|
||
) -> list[Record]:
|
||
"""Batch-insert multiple records (P2 #19: eliminates per-record INSERT).
|
||
|
||
Returns the created records in insertion order.
|
||
"""
|
||
if not records_values:
|
||
return []
|
||
async with self._session_factory() as session:
|
||
rows = [
|
||
{"id": _uuid_str(), "table_id": table_id, "values": vals} for vals in records_values
|
||
]
|
||
stmt = insert(RecordModel).values(rows).returning(RecordModel)
|
||
result = await session.execute(stmt)
|
||
entities = result.scalars().all()
|
||
await session.commit()
|
||
return [Record.model_validate(e) for e in entities]
|
||
|
||
async def get_record(self, record_id: str) -> Record | None:
|
||
"""Get a record by ID."""
|
||
async with self._session_factory() as session:
|
||
stmt = select(RecordModel).where(RecordModel.id == record_id)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
return Record.model_validate(entity) if entity else None
|
||
|
||
async def list_records(
|
||
self,
|
||
table_id: str,
|
||
cursor: str | None = None,
|
||
limit: int = 50,
|
||
) -> tuple[list[Record], str | None]:
|
||
"""List records in a table with cursor-based pagination.
|
||
|
||
Returns ``(records, next_cursor)``. ``next_cursor`` is ``None`` when
|
||
there are no more records.
|
||
"""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
select(RecordModel)
|
||
.where(RecordModel.table_id == table_id)
|
||
.order_by(RecordModel.id)
|
||
.limit(limit + 1) # fetch one extra to check if there's a next page
|
||
)
|
||
if cursor:
|
||
stmt = stmt.where(RecordModel.id > cursor)
|
||
|
||
result = await session.execute(stmt)
|
||
entities = result.scalars().all()
|
||
|
||
if len(entities) > limit:
|
||
next_cursor = entities[limit - 1].id
|
||
entities = entities[:limit]
|
||
else:
|
||
next_cursor = None
|
||
|
||
return [Record.model_validate(e) for e in entities], next_cursor
|
||
|
||
async def update_record_values(
|
||
self, record_id: str, values: dict[str, object]
|
||
) -> Record | None:
|
||
"""Update a record's values (full replace)."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
update(RecordModel)
|
||
.where(RecordModel.id == record_id)
|
||
.values(values=values)
|
||
.returning(RecordModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
await session.commit()
|
||
return Record.model_validate(entity) if entity else None
|
||
|
||
async def delete_record(self, record_id: str) -> bool:
|
||
"""Delete a record."""
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(delete(RecordModel).where(RecordModel.id == record_id))
|
||
await session.commit()
|
||
return result.rowcount > 0
|
||
|
||
async def delete_records_by_table(self, table_id: str) -> int:
|
||
"""Delete all records in a table. Returns count deleted."""
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(
|
||
delete(RecordModel).where(RecordModel.table_id == table_id)
|
||
)
|
||
await session.commit()
|
||
return result.rowcount
|
||
|
||
# ── Views ───────────────────────────────────────────────
|
||
|
||
async def create_view(
|
||
self,
|
||
table_id: str,
|
||
name: str,
|
||
view_type: ViewType = ViewType.grid,
|
||
config: dict[str, object] | None = None,
|
||
) -> View:
|
||
"""Create a new view."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
insert(ViewModel)
|
||
.values(
|
||
id=_uuid_str(),
|
||
table_id=table_id,
|
||
name=name,
|
||
view_type=view_type.value,
|
||
config=config or {},
|
||
)
|
||
.returning(ViewModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalar_one()
|
||
await session.commit()
|
||
return View.model_validate(entity)
|
||
|
||
async def get_view(self, view_id: str) -> View | None:
|
||
"""Get a single view by ID (used for ownership checks)."""
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(select(ViewModel).where(ViewModel.id == view_id))
|
||
entity = result.scalars().first()
|
||
return View.model_validate(entity) if entity else None
|
||
|
||
async def list_views(self, table_id: str) -> list[View]:
|
||
"""List all views in a table."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
select(ViewModel)
|
||
.where(ViewModel.table_id == table_id)
|
||
.order_by(ViewModel.created_at)
|
||
)
|
||
result = await session.execute(stmt)
|
||
return [View.model_validate(e) for e in result.scalars().all()]
|
||
|
||
async def update_view(self, view_id: str, **kwargs: object) -> View | None:
|
||
"""Update a view's attributes."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
update(ViewModel)
|
||
.where(ViewModel.id == view_id)
|
||
.values(**kwargs)
|
||
.returning(ViewModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
await session.commit()
|
||
return View.model_validate(entity) if entity else None
|
||
|
||
async def delete_view_if_not_last(self, view_id: str, table_id: str) -> bool:
|
||
"""Atomically delete a view only if it's not the last one in its table (DR-4 TOCTOU fix).
|
||
|
||
Uses ``SELECT ... FOR UPDATE`` in a single transaction to lock sibling
|
||
rows for the duration of the count + delete, preventing concurrent
|
||
``delete_view`` calls from racing past the last-view check and leaving
|
||
the table with zero views.
|
||
|
||
Returns:
|
||
True if the view was deleted.
|
||
False if NOT deleted because it would be the last view (or the
|
||
view doesn't belong to the given table_id).
|
||
"""
|
||
async with self._session_factory() as session:
|
||
# Lock all sibling views — concurrent delete_view calls on the
|
||
# same table block here until this transaction commits.
|
||
lock_sql = text(
|
||
"SELECT id FROM bitable.bitable_views WHERE table_id = :table_id FOR UPDATE"
|
||
)
|
||
result = await session.execute(lock_sql, {"table_id": table_id})
|
||
sibling_ids = {str(row[0]) for row in result.fetchall()}
|
||
if len(sibling_ids) <= 1:
|
||
return False # Last view — refuse to delete
|
||
if view_id not in sibling_ids:
|
||
return False # View doesn't belong to this table_id
|
||
result = await session.execute(delete(ViewModel).where(ViewModel.id == view_id))
|
||
await session.commit()
|
||
return result.rowcount > 0
|
||
|
||
# ── Recalc Queue ────────────────────────────────────────
|
||
|
||
async def enqueue_recalc(
|
||
self, table_id: str, record_id: str, field_id: str
|
||
) -> RecalcTask | None:
|
||
"""Enqueue a recalc task. Returns None if duplicate (ON CONFLICT DO NOTHING)."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
pg_insert(RecalcQueueModel)
|
||
.values(
|
||
id=_uuid_str(),
|
||
table_id=table_id,
|
||
record_id=record_id,
|
||
field_id=field_id,
|
||
status=RecalcStatus.pending.value,
|
||
)
|
||
.on_conflict_do_nothing(constraint="uq_recalc_record_field")
|
||
.returning(RecalcQueueModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
entity = result.scalars().first()
|
||
await session.commit()
|
||
return RecalcTask.model_validate(entity) if entity else None
|
||
|
||
async def get_pending_recalc_tasks(self, limit: int = 50) -> list[RecalcTask]:
|
||
"""Get pending recalc tasks ordered by queue time.
|
||
|
||
Deprecated for concurrent workers — use :meth:`claim_recalc_tasks`
|
||
for atomic claim. Kept for introspection/tests.
|
||
"""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
select(RecalcQueueModel)
|
||
.where(RecalcQueueModel.status == RecalcStatus.pending.value)
|
||
.order_by(RecalcQueueModel.queued_at)
|
||
.limit(limit)
|
||
)
|
||
result = await session.execute(stmt)
|
||
return [RecalcTask.model_validate(e) for e in result.scalars().all()]
|
||
|
||
async def claim_recalc_tasks(self, limit: int = 10) -> list[RecalcTask]:
|
||
"""Atomically claim up to ``limit`` pending tasks (P1 #6).
|
||
|
||
Uses ``FOR UPDATE SKIP LOCKED`` so multiple workers never grab the
|
||
same task. Each claimed task is set to ``calculating`` in the same
|
||
transaction and returned. Callers process the returned tasks and
|
||
then call :meth:`update_recalc_status` with ``done``/``error``.
|
||
"""
|
||
# ponytail: PostgreSQL-specific FOR UPDATE SKIP LOCKED. Ceiling: this
|
||
# binds the repository to PG; upgrade path is a dialect check.
|
||
async with self._session_factory() as session:
|
||
# Subselect pending tasks with row-level locks, skipping locked rows.
|
||
subq = (
|
||
select(RecalcQueueModel.id)
|
||
.where(RecalcQueueModel.status == RecalcStatus.pending.value)
|
||
.order_by(RecalcQueueModel.queued_at)
|
||
.limit(limit)
|
||
.with_for_update(skip_locked=True)
|
||
).subquery()
|
||
|
||
stmt = (
|
||
update(RecalcQueueModel)
|
||
.where(RecalcQueueModel.id.in_(select(subq.c.id)))
|
||
.values(status=RecalcStatus.calculating.value)
|
||
.returning(RecalcQueueModel)
|
||
)
|
||
result = await session.execute(stmt)
|
||
await session.commit()
|
||
return [RecalcTask.model_validate(e) for e in result.scalars().all()]
|
||
|
||
async def update_recalc_status(
|
||
self,
|
||
task_id: str,
|
||
status: RecalcStatus,
|
||
error_message: str | None = None,
|
||
) -> None:
|
||
"""Update a recalc task's status."""
|
||
async with self._session_factory() as session:
|
||
kwargs: dict[str, object] = {"status": status.value}
|
||
if error_message is not None:
|
||
kwargs["error_message"] = error_message
|
||
if status in (RecalcStatus.done, RecalcStatus.error):
|
||
kwargs["completed_at"] = func.now()
|
||
stmt = update(RecalcQueueModel).where(RecalcQueueModel.id == task_id).values(**kwargs)
|
||
await session.execute(stmt)
|
||
await session.commit()
|
||
|
||
async def reset_stale_recalc_tasks(self, stale_threshold: float = 600.0) -> int:
|
||
"""Reset 'calculating' tasks older than ``stale_threshold`` seconds back to 'pending'.
|
||
|
||
Only tasks whose ``queued_at`` is older than ``now - stale_threshold``
|
||
are reset — this avoids resetting tasks that a live worker is currently
|
||
processing (crash recovery, not active-task interruption). Returns count.
|
||
"""
|
||
# ponytail: stale_threshold is the max age a calculating task should
|
||
# reach. 600s default = 10min; a healthy worker finishes in milliseconds.
|
||
cutoff = datetime.now(timezone.utc) - timedelta(seconds=stale_threshold)
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
update(RecalcQueueModel)
|
||
.where(RecalcQueueModel.status == RecalcStatus.calculating.value)
|
||
.where(RecalcQueueModel.queued_at < cutoff)
|
||
.values(status=RecalcStatus.pending.value)
|
||
)
|
||
result = await session.execute(stmt)
|
||
await session.commit()
|
||
return result.rowcount
|
||
|
||
# ── Upsert (KTD8: jsonb_set per agent field) ───────────
|
||
|
||
async def find_record_by_pk(
|
||
self, table_id: str, pk_field_id: str, pk_value: str
|
||
) -> RecordModel | None:
|
||
"""Find a record by primary key value (JSONB key lookup)."""
|
||
async with self._session_factory() as session:
|
||
# ponytail: field_id is a system-generated UUID, safe to interpolate.
|
||
# pk_value is parameterized.
|
||
sql = text(
|
||
"SELECT * FROM bitable.bitable_records "
|
||
"WHERE table_id = :table_id AND values->>:pk_path = :pk_value LIMIT 1"
|
||
)
|
||
result = await session.execute(
|
||
sql, {"table_id": table_id, "pk_path": pk_field_id, "pk_value": pk_value}
|
||
)
|
||
row = result.fetchone()
|
||
return RecordModel(**row._mapping) if row else None
|
||
|
||
async def find_records_by_pk_batch(
|
||
self, table_id: str, pk_field_id: str, pk_values: list[str]
|
||
) -> dict[str, RecordModel]:
|
||
"""Batch-find records by PK values (P1 #14: eliminates N+1 queries).
|
||
|
||
Returns a dict mapping pk_value (str) → RecordModel.
|
||
"""
|
||
if not pk_values:
|
||
return {}
|
||
async with self._session_factory() as session:
|
||
# ponytail: field_id is a system UUID, safe to interpolate.
|
||
# pk_values are parameterized via ANY(:values).
|
||
sql = text(
|
||
"SELECT * FROM bitable.bitable_records "
|
||
"WHERE table_id = :table_id AND values->>:pk_path = ANY(:pk_values)"
|
||
)
|
||
result = await session.execute(
|
||
sql,
|
||
{
|
||
"table_id": table_id,
|
||
"pk_path": pk_field_id,
|
||
"pk_values": pk_values,
|
||
},
|
||
)
|
||
rows = result.fetchall()
|
||
result_map: dict[str, RecordModel] = {}
|
||
for r in rows:
|
||
values = r._mapping["values"]
|
||
if isinstance(values, str):
|
||
import json as _json
|
||
|
||
values = _json.loads(values)
|
||
pk_val = values.get(pk_field_id) if isinstance(values, dict) else None
|
||
if pk_val is not None:
|
||
result_map[str(pk_val)] = RecordModel(**r._mapping)
|
||
return result_map
|
||
|
||
async def upsert_record_agent_fields(
|
||
self, record_id: str, agent_field_values: dict[str, object]
|
||
) -> None:
|
||
"""Update agent-owned fields using jsonb_set (KTD8).
|
||
|
||
Chains jsonb_set calls so user-owned fields are never touched.
|
||
"""
|
||
if not agent_field_values:
|
||
return
|
||
import json
|
||
|
||
# U3/DR-1: 校验所有 field_id 是 UUID 格式后再插值到 jsonb_set path（防御性，避免 SQL 结构注入）
|
||
for fid in agent_field_values:
|
||
_validate_field_id(fid, context="upsert_record_agent_fields")
|
||
|
||
async with self._session_factory() as session:
|
||
# Build nested jsonb_set: jsonb_set(jsonb_set(values, '{f1}', CAST(:v0 AS jsonb), true), ...)
|
||
# ponytail: field_ids validated above as UUID format, safe to interpolate into path literals.
|
||
# Use CAST(:param AS jsonb) instead of :param::jsonb — asyncpg dialect
|
||
# misparses the `::` as part of the param name.
|
||
inner = "values"
|
||
params: dict[str, object] = {"record_id": record_id}
|
||
for i, (field_id, value) in enumerate(agent_field_values.items()):
|
||
param_key = f"v{i}"
|
||
inner = f"jsonb_set({inner}, '{{{field_id}}}', CAST(:{param_key} AS jsonb), true)"
|
||
params[param_key] = json.dumps(value)
|
||
sql = text(f"UPDATE bitable.bitable_records SET values = {inner} WHERE id = :record_id")
|
||
await session.execute(sql, params)
|
||
await session.commit()
|
||
|
||
# ── View-filtered record listing (KTD9) ────────────────
|
||
|
||
async def list_records_filtered(
|
||
self,
|
||
table_id: str,
|
||
filters: list[dict[str, object]] | None = None,
|
||
sorts: list[dict[str, object]] | None = None,
|
||
cursor: str | None = None,
|
||
limit: int = 50,
|
||
) -> tuple[list[Record], str | None]:
|
||
"""List records with view filters/sorts + cursor pagination.
|
||
|
||
filters: [{"field_id": "...", "op": "eq|ne|gt|lt|contains|is_empty", "value": ...}]
|
||
sorts: [{"field_id": "...", "direction": "asc|desc"}]
|
||
|
||
Cursor is a base64-encoded JSON of ``{"id": ..., "sv": [sort_val, ...]}``.
|
||
When no sorts are given, the cursor degrades to id-only pagination.
|
||
With sorts, a row-value comparison ``(sort_cols..., id) > (cursor_vals...)``
|
||
ensures stable pagination across non-id orderings (P1 #11).
|
||
"""
|
||
import base64
|
||
import json as _json
|
||
|
||
# U3/DR-1: 校验 filters/sorts 中的 field_id 是 UUID 格式后再插值到 SQL 结构
|
||
if filters:
|
||
for f in filters:
|
||
_validate_field_id(f["field_id"], context="list_records_filtered.filter")
|
||
if sorts:
|
||
for s in sorts:
|
||
_validate_field_id(s["field_id"], context="list_records_filtered.sort")
|
||
|
||
async with self._session_factory() as session:
|
||
# Build raw SQL with JSONB filter/sort translation.
|
||
# ponytail: field_ids validated above as UUID format, safe to interpolate into SQL structure.
|
||
where_clauses = ["table_id = :table_id"]
|
||
params: dict[str, object] = {"table_id": table_id}
|
||
|
||
if filters:
|
||
for i, f in enumerate(filters):
|
||
fid = f["field_id"]
|
||
op = f["op"]
|
||
val = f.get("value")
|
||
param_name = f"fv{i}"
|
||
if op == "is_empty":
|
||
where_clauses.append(f"(values->>'{fid}') IS NULL OR values->>'{fid}' = ''")
|
||
elif op == "eq":
|
||
where_clauses.append(f"values->>'{fid}' = :{param_name}")
|
||
params[param_name] = str(val)
|
||
elif op == "ne":
|
||
where_clauses.append(f"values->>'{fid}' != :{param_name}")
|
||
params[param_name] = str(val)
|
||
elif op == "contains":
|
||
where_clauses.append(f"values->>'{fid}' LIKE :{param_name}")
|
||
params[param_name] = f"%{val}%"
|
||
elif op in ("gt", "lt", "gte", "lte"):
|
||
op_map = {"gt": ">", "lt": "<", "gte": ">=", "lte": "<="}
|
||
where_clauses.append(
|
||
f"CAST(values->>'{fid}' AS NUMERIC) {op_map[op]} :{param_name}"
|
||
)
|
||
params[param_name] = val
|
||
|
||
# Build sort expressions and cursor (P1 #11: composite cursor).
|
||
sort_exprs: list[str] = [] # SQL expressions for ORDER BY
|
||
sort_dirs: list[str] = [] # "ASC" or "DESC" per sort column
|
||
if sorts:
|
||
for s in sorts:
|
||
fid = s["field_id"]
|
||
direction = "ASC" if s.get("direction", "asc") == "asc" else "DESC"
|
||
sort_exprs.append(f"values->>'{fid}'")
|
||
sort_dirs.append(direction)
|
||
# Always append id as the final tiebreaker for stable ordering.
|
||
sort_exprs.append("id")
|
||
sort_dirs.append(
|
||
"ASC" if not sorts or sorts[0].get("direction", "asc") == "asc" else "DESC"
|
||
)
|
||
|
||
order_by = ", ".join(f"{expr} {dir_}" for expr, dir_ in zip(sort_exprs, sort_dirs))
|
||
|
||
# Cursor: composite row-value comparison.
|
||
if cursor:
|
||
try:
|
||
decoded = _json.loads(base64.b64decode(cursor).decode("utf-8"))
|
||
cursor_id = decoded["id"]
|
||
cursor_sort_vals = decoded.get("sv", [])
|
||
except (ValueError, KeyError) as e:
|
||
raise ValueError(f"Invalid cursor: {e}") from e
|
||
|
||
if not cursor_sort_vals:
|
||
# No sorts — simple id comparison.
|
||
where_clauses.append("id > :cursor_id")
|
||
params["cursor_id"] = cursor_id
|
||
else:
|
||
# Row-value comparison: (col1, col2, ..., id) > (v1, v2, ..., cursor_id)
|
||
# For DESC columns, invert the comparison direction per-column
|
||
# by negating the value or using < instead of >.
|
||
# PostgreSQL row values: (a,b) > (x,y) means a>x OR (a=x AND b>y).
|
||
# For DESC, we want a<x OR (a=x AND b<x), so we use < and the
|
||
# ORDER BY already has DESC.
|
||
# Build the LHS and RHS of the row comparison.
|
||
lhs_parts: list[str] = []
|
||
rhs_params: list[str] = []
|
||
for i, expr in enumerate(sort_exprs):
|
||
lhs_parts.append(expr)
|
||
param_name = f"csv{i}"
|
||
rhs_params.append(f":{param_name}")
|
||
params[param_name] = cursor_sort_vals[i]
|
||
# Last element is id.
|
||
lhs_parts.append("id")
|
||
params["cursor_id"] = cursor_id
|
||
rhs_params.append(":cursor_id")
|
||
|
||
# Determine comparison direction: if first sort is DESC, use <.
|
||
# Mixed-direction sorts are rare; for correctness with mixed
|
||
# directions, we'd need per-column operators. For now, use the
|
||
# first sort's direction (covers the common single-sort case).
|
||
# ponytail: ceiling — mixed ASC/DESC sorts use first direction only.
|
||
comp_op = "<" if sort_dirs[0] == "DESC" else ">"
|
||
lhs = f"({', '.join(lhs_parts)})"
|
||
rhs = f"({', '.join(rhs_params)})"
|
||
where_clauses.append(f"{lhs} {comp_op} {rhs}")
|
||
|
||
where_sql = " AND ".join(where_clauses)
|
||
sql = text(
|
||
f"SELECT * FROM bitable.bitable_records WHERE {where_sql} "
|
||
f"ORDER BY {order_by} LIMIT :limit"
|
||
)
|
||
params["limit"] = limit + 1
|
||
|
||
result = await session.execute(sql, params)
|
||
rows = result.fetchall()
|
||
|
||
if len(rows) > limit:
|
||
last_row = rows[limit - 1]
|
||
last_mapping = last_row._mapping
|
||
# Build composite cursor from sort values + id.
|
||
# Sort values are extracted as text to match `values->>'fid'` expressions.
|
||
sv: list[object] = []
|
||
last_values = last_mapping.get("values")
|
||
if isinstance(last_values, str):
|
||
# asyncpg may return JSONB as str in raw text() queries.
|
||
try:
|
||
last_values = _json.loads(last_values)
|
||
except (ValueError, TypeError):
|
||
last_values = {}
|
||
if sorts and isinstance(last_values, dict):
|
||
for s in sorts:
|
||
fid = s["field_id"]
|
||
val = last_values.get(fid)
|
||
sv.append(str(val) if val is not None else None)
|
||
cursor_data = {"id": last_mapping["id"], "sv": sv}
|
||
next_cursor = base64.b64encode(_json.dumps(cursor_data).encode("utf-8")).decode(
|
||
"utf-8"
|
||
)
|
||
rows = rows[:limit]
|
||
else:
|
||
next_cursor = None
|
||
|
||
return [Record.model_validate(RecordModel(**r._mapping)) for r in rows], next_cursor
|
||
|
||
# ── Field deletion dependency queries ──────────────────
|
||
|
||
async def find_formula_fields_referencing(
|
||
self, table_id: str, target_field_id: str
|
||
) -> list[Field]:
|
||
"""Find formula fields that reference target_field_id in their formula_expr."""
|
||
async with self._session_factory() as session:
|
||
# JSONB config->>'formula_expr' LIKE '%{target_field_id}%'
|
||
sql = text(
|
||
"SELECT * FROM bitable.bitable_fields "
|
||
"WHERE table_id = :table_id AND field_type = 'formula' "
|
||
"AND config->>'formula_expr' LIKE :pattern"
|
||
)
|
||
result = await session.execute(
|
||
sql,
|
||
{"table_id": table_id, "pattern": f"%{target_field_id}%"},
|
||
)
|
||
return [Field.model_validate(FieldModel(**r._mapping)) for r in result.fetchall()]
|
||
|
||
async def find_views_referencing_field(self, table_id: str, field_id: str) -> list[View]:
|
||
"""Find views whose config (filters/sorts/hidden_fields) references field_id."""
|
||
async with self._session_factory() as session:
|
||
# Check if field_id appears anywhere in config JSONB
|
||
sql = text(
|
||
"SELECT * FROM bitable.bitable_views "
|
||
"WHERE table_id = :table_id AND config::text LIKE :pattern"
|
||
)
|
||
result = await session.execute(
|
||
sql,
|
||
{"table_id": table_id, "pattern": f"%{field_id}%"},
|
||
)
|
||
return [View.model_validate(ViewModel(**r._mapping)) for r in result.fetchall()]
|
||
|
||
async def remove_field_from_records(self, table_id: str, field_id: str) -> int:
|
||
"""Remove a field key from all records' JSONB values (force delete cleanup)."""
|
||
async with self._session_factory() as session:
|
||
sql = text(
|
||
"UPDATE bitable.bitable_records SET values = values - :field_path "
|
||
"WHERE table_id = :table_id"
|
||
)
|
||
result = await session.execute(sql, {"field_path": field_id, "table_id": table_id})
|
||
await session.commit()
|
||
return result.rowcount
|
||
|
||
# ── Recalc support (U3) ────────────────────────────────
|
||
|
||
async def get_column_values(self, table_id: str, field_id: str) -> list[object]:
|
||
"""Get all values for a field across all records in a table (for aggregates).
|
||
|
||
Returns a list of values (preserving order by record id). Missing values
|
||
are included as None so aggregate functions can skip them.
|
||
"""
|
||
async with self._session_factory() as session:
|
||
sql = text(
|
||
"SELECT values->:field_id AS val FROM bitable.bitable_records "
|
||
"WHERE table_id = :table_id ORDER BY id"
|
||
)
|
||
result = await session.execute(sql, {"field_id": field_id, "table_id": table_id})
|
||
return [row[0] for row in result.fetchall()]
|
||
|
||
async def set_formula_value(self, record_id: str, field_id: str, value: object) -> None:
|
||
"""Set a single formula field value in a record's JSONB (jsonb_set)."""
|
||
import json
|
||
|
||
async with self._session_factory() as session:
|
||
# ponytail: field_id is a system UUID, safe to interpolate into path literal.
|
||
# Use CAST AS jsonb instead of ::jsonb (asyncpg dialect misparses ::).
|
||
sql = text(
|
||
f"UPDATE bitable.bitable_records "
|
||
f"SET values = jsonb_set(values, '{{{field_id}}}', CAST(:val AS jsonb), true) "
|
||
f"WHERE id = :record_id"
|
||
)
|
||
await session.execute(sql, {"val": json.dumps(value), "record_id": record_id})
|
||
await session.commit()
|
||
|
||
async def get_all_records(self, table_id: str) -> list[Record]:
|
||
"""Get all records in a table (for column value extraction in recalc)."""
|
||
async with self._session_factory() as session:
|
||
stmt = (
|
||
select(RecordModel).where(RecordModel.table_id == table_id).order_by(RecordModel.id)
|
||
)
|
||
result = await session.execute(stmt)
|
||
return [Record.model_validate(e) for e in result.scalars().all()]
|