408 lines
14 KiB
Python
408 lines
14 KiB
Python
"""PostgreSQL schema and initialization for the bitable subsystem.
|
|
|
|
Uses an independent ``bitable`` schema within the shared PostgreSQL instance.
|
|
Follows the lazy-init + lock pattern from ``evolution/pg_store.py`` and the
|
|
``_SCHEMA_VERSION`` migration pattern from ``server/auth/models.py``.
|
|
|
|
Schema versioning
|
|
-----------------
|
|
:data:`_SCHEMA_VERSION` tracks the current bitable DB schema. The
|
|
``bitable_meta`` table stores the version so subsequent restarts skip
|
|
already-applied migrations.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import uuid as _uuid
|
|
from datetime import datetime, timezone
|
|
from types import TracebackType
|
|
|
|
from sqlalchemy import (
|
|
Column,
|
|
DateTime,
|
|
Index,
|
|
String,
|
|
Text,
|
|
UniqueConstraint,
|
|
text,
|
|
)
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Current schema version — bump when adding migrations.
|
|
# V1: initial schema (tables/fields/records/views/recalc/meta).
|
|
# V2: add bitable_files table + file_id column on bitable_tables (R1).
|
|
_SCHEMA_VERSION = 2
|
|
|
|
_META_SCHEMA_VERSION_KEY = "schema_version"
|
|
|
|
|
|
def _utcnow() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def _uuid_str() -> str:
|
|
return str(_uuid.uuid4())
|
|
|
|
|
|
class BitableBase(DeclarativeBase):
|
|
"""Declarative base for bitable ORM models (independent schema)."""
|
|
|
|
|
|
class FileModel(BitableBase):
|
|
"""ORM model for ``bitable.bitable_files`` — top-level container (R1)."""
|
|
|
|
__tablename__ = "bitable_files"
|
|
__table_args__ = (
|
|
Index("ix_bitable_files_owner", "owner_user_id"),
|
|
{"schema": "bitable"},
|
|
)
|
|
|
|
id = Column(String, primary_key=True, default=_uuid_str)
|
|
name = Column(String, nullable=False)
|
|
icon = Column(String, default="table")
|
|
description = Column(Text, default="")
|
|
owner_user_id = Column(String, nullable=True)
|
|
created_at = Column(DateTime(timezone=True), default=_utcnow)
|
|
updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)
|
|
|
|
|
|
class TableModel(BitableBase):
|
|
"""ORM model for ``bitable.bitable_tables``."""
|
|
|
|
__tablename__ = "bitable_tables"
|
|
__table_args__ = (
|
|
Index("ix_bitable_tables_file_id", "file_id"),
|
|
{"schema": "bitable"},
|
|
)
|
|
|
|
id = Column(String, primary_key=True, default=_uuid_str)
|
|
file_id = Column(String, nullable=True)
|
|
name = Column(String, nullable=False)
|
|
description = Column(Text, default="")
|
|
primary_key_field_id = Column(String, nullable=True)
|
|
owner_user_id = Column(String, nullable=True)
|
|
created_at = Column(DateTime(timezone=True), default=_utcnow)
|
|
updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)
|
|
|
|
|
|
class FieldModel(BitableBase):
|
|
"""ORM model for ``bitable.bitable_fields``."""
|
|
|
|
__tablename__ = "bitable_fields"
|
|
__table_args__ = (
|
|
Index("ix_bitable_fields_table_id", "table_id"),
|
|
{"schema": "bitable"},
|
|
)
|
|
|
|
id = Column(String, primary_key=True, default=_uuid_str)
|
|
table_id = Column(String, nullable=False)
|
|
name = Column(String, nullable=False)
|
|
field_type = Column(String, nullable=False)
|
|
config = Column(JSONB, default=dict)
|
|
owner = Column(String, default="user")
|
|
created_at = Column(DateTime(timezone=True), default=_utcnow)
|
|
|
|
|
|
class RecordModel(BitableBase):
|
|
"""ORM model for ``bitable.bitable_records``.
|
|
|
|
``values`` is JSONB mapping ``{field_id: value}``. A GIN index supports
|
|
efficient JSONB key existence checks. A unique expression index on
|
|
``(table_id, values->>pk_field_id)`` enforces primary-key uniqueness
|
|
(created dynamically in ``_apply_v1_schema`` because the pk field id is
|
|
per-table, not a fixed column).
|
|
"""
|
|
|
|
__tablename__ = "bitable_records"
|
|
__table_args__ = (
|
|
Index("ix_bitable_records_table_id", "table_id"),
|
|
Index("ix_bitable_records_values_gin", "values", postgresql_using="gin"),
|
|
{"schema": "bitable"},
|
|
)
|
|
|
|
id = Column(String, primary_key=True, default=_uuid_str)
|
|
table_id = Column(String, nullable=False)
|
|
values = Column(JSONB, default=dict)
|
|
created_at = Column(DateTime(timezone=True), default=_utcnow)
|
|
updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)
|
|
|
|
|
|
class ViewModel(BitableBase):
|
|
"""ORM model for ``bitable.bitable_views``."""
|
|
|
|
__tablename__ = "bitable_views"
|
|
__table_args__ = (
|
|
Index("ix_bitable_views_table_id", "table_id"),
|
|
{"schema": "bitable"},
|
|
)
|
|
|
|
id = Column(String, primary_key=True, default=_uuid_str)
|
|
table_id = Column(String, nullable=False)
|
|
name = Column(String, nullable=False)
|
|
view_type = Column(String, default="grid")
|
|
config = Column(JSONB, default=dict)
|
|
created_at = Column(DateTime(timezone=True), default=_utcnow)
|
|
|
|
|
|
class RecalcQueueModel(BitableBase):
|
|
"""ORM model for ``bitable.bitable_recalc_queue``.
|
|
|
|
The ``(record_id, field_id)`` unique index prevents duplicate enqueues.
|
|
The ``(status, queued_at)`` index supports efficient worker consumption.
|
|
"""
|
|
|
|
__tablename__ = "bitable_recalc_queue"
|
|
__table_args__ = (
|
|
UniqueConstraint("record_id", "field_id", name="uq_recalc_record_field"),
|
|
Index("ix_recalc_status_queued", "status", "queued_at"),
|
|
{"schema": "bitable"},
|
|
)
|
|
|
|
id = Column(String, primary_key=True, default=_uuid_str)
|
|
table_id = Column(String, nullable=False)
|
|
record_id = Column(String, nullable=False)
|
|
field_id = Column(String, nullable=False)
|
|
status = Column(String, default="pending")
|
|
error_message = Column(Text, nullable=True)
|
|
queued_at = Column(DateTime(timezone=True), default=_utcnow)
|
|
completed_at = Column(DateTime(timezone=True), nullable=True)
|
|
|
|
|
|
class MetaModel(BitableBase):
|
|
"""ORM model for ``bitable.bitable_meta`` — schema version tracking."""
|
|
|
|
__tablename__ = "bitable_meta"
|
|
__table_args__ = {"schema": "bitable"}
|
|
|
|
key = Column(String, primary_key=True)
|
|
value = Column(String, nullable=False)
|
|
updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schema is created via BitableBase.metadata.create_all (see BitableDB.init).
|
|
# The ORM models above are the single source of truth for table/index DDL.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _apply_v2_migration(conn: object) -> None:
|
|
"""V2 migration: create ``bitable_files`` table + add ``file_id`` to tables.
|
|
|
|
Idempotent — safe to call on fresh installs (``create_all`` already made
|
|
the files table and file_id column) and on V1 upgrades (ALTERs the
|
|
existing tables). The CREATE TABLE IF NOT EXISTS + ADD COLUMN IF NOT
|
|
EXISTS guards make both paths work.
|
|
|
|
Defensive: V1 tables without a ``file_id`` are left with ``file_id = NULL``
|
|
(orphaned). The service layer treats NULL file_id as "no parent file"
|
|
and continues to work — they're reachable via the deprecated
|
|
``/tables`` endpoint. The plan notes there is no existing production
|
|
data to migrate (verifier-confirmed), so this is purely defensive.
|
|
"""
|
|
# bitable_files table — created by create_all on fresh installs, but
|
|
# we recreate via raw SQL so V1→V2 upgrades (no create_all re-run for
|
|
# existing tables) also get it.
|
|
await conn.execute(
|
|
text(
|
|
"CREATE TABLE IF NOT EXISTS bitable.bitable_files ("
|
|
" id VARCHAR PRIMARY KEY,"
|
|
" name VARCHAR NOT NULL,"
|
|
" icon VARCHAR DEFAULT 'table',"
|
|
" description TEXT DEFAULT '',"
|
|
" owner_user_id VARCHAR,"
|
|
" created_at TIMESTAMPTZ DEFAULT NOW(),"
|
|
" updated_at TIMESTAMPTZ DEFAULT NOW()"
|
|
")"
|
|
)
|
|
)
|
|
await conn.execute(
|
|
text(
|
|
"CREATE INDEX IF NOT EXISTS ix_bitable_files_owner "
|
|
"ON bitable.bitable_files (owner_user_id)"
|
|
)
|
|
)
|
|
|
|
# Add file_id column to bitable_tables (idempotent).
|
|
# ponytail: ADD COLUMN IF NOT EXISTS is PostgreSQL 9.6+ — safe here.
|
|
await conn.execute(
|
|
text("ALTER TABLE bitable.bitable_tables ADD COLUMN IF NOT EXISTS file_id VARCHAR")
|
|
)
|
|
await conn.execute(
|
|
text(
|
|
"CREATE INDEX IF NOT EXISTS ix_bitable_tables_file_id "
|
|
"ON bitable.bitable_tables (file_id)"
|
|
)
|
|
)
|
|
logger.info("applied V2 migration: bitable_files table + file_id column")
|
|
|
|
|
|
def _resolve_database_url(database_url: str | None = None) -> str | None:
|
|
"""Resolve PostgreSQL connection URL.
|
|
|
|
Priority: explicit arg > env ``DATABASE_URL`` > ``AGENTKIT_DATABASE_URL``.
|
|
"""
|
|
if database_url:
|
|
return database_url
|
|
return os.environ.get("DATABASE_URL") or os.environ.get("AGENTKIT_DATABASE_URL")
|
|
|
|
|
|
class BitableDB:
|
|
"""PostgreSQL connection manager for bitable (lazy init + lock).
|
|
|
|
Usage::
|
|
|
|
db = BitableDB(database_url="postgresql+asyncpg://...")
|
|
await db.init()
|
|
# ... use db.engine / db.session_factory ...
|
|
await db.close()
|
|
"""
|
|
|
|
def __init__(self, database_url: str | None = None) -> None:
|
|
self._database_url = database_url or _resolve_database_url()
|
|
self._engine: object | None = None
|
|
self._session_factory: object | None = None
|
|
self._initialized = False
|
|
self._init_lock = asyncio.Lock()
|
|
|
|
@property
|
|
def database_url(self) -> str | None:
|
|
return self._database_url
|
|
|
|
@property
|
|
def engine(self) -> object | None:
|
|
return self._engine
|
|
|
|
@property
|
|
def session_factory(self) -> object | None:
|
|
return self._session_factory
|
|
|
|
@property
|
|
def is_initialized(self) -> bool:
|
|
return self._initialized
|
|
|
|
async def _ensure_initialized(self) -> None:
|
|
"""Lazy-init async engine and session factory (with lock)."""
|
|
if self._initialized:
|
|
return
|
|
|
|
async with self._init_lock:
|
|
if self._initialized:
|
|
return
|
|
|
|
if not self._database_url:
|
|
raise RuntimeError(
|
|
"No database URL configured for bitable. "
|
|
"Set DATABASE_URL or AGENTKIT_DATABASE_URL env var."
|
|
)
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
self._engine = create_async_engine(self._database_url, echo=False)
|
|
self._session_factory = sessionmaker(
|
|
self._engine, class_=AsyncSession, expire_on_commit=False
|
|
)
|
|
self._initialized = True
|
|
|
|
async def init(self) -> None:
|
|
"""Initialize engine + create schema and tables (idempotent).
|
|
|
|
Runs schema migrations based on ``_SCHEMA_VERSION`` stored in
|
|
``bitable_meta``. Safe to call on every startup.
|
|
"""
|
|
await self._ensure_initialized()
|
|
|
|
async with self._engine.begin() as conn:
|
|
# 1. Create the bitable schema (idempotent) — must precede
|
|
# metadata.create_all since tables live in this schema.
|
|
await conn.execute(text("CREATE SCHEMA IF NOT EXISTS bitable"))
|
|
|
|
# 2. Create all tables/indexes/constraints from ORM metadata.
|
|
# Uses run_sync because asyncpg doesn't support multi-statement
|
|
# text() execution; metadata.create_all emits one DDL statement
|
|
# per table and handles schema-qualified names + GIN indexes.
|
|
await conn.run_sync(BitableBase.metadata.create_all)
|
|
|
|
# 3. Check current schema version
|
|
result = await conn.execute(
|
|
text("SELECT value FROM bitable.bitable_meta WHERE key = :key"),
|
|
{"key": _META_SCHEMA_VERSION_KEY},
|
|
)
|
|
row = result.fetchone()
|
|
current_version = int(row[0]) if row else 0
|
|
|
|
# 4. Apply migrations if needed (future versions add elif blocks here)
|
|
if current_version < _SCHEMA_VERSION:
|
|
# V1: initial schema (already created above via create_all)
|
|
# V2: file layer (bitable_files table + file_id column on tables)
|
|
if current_version < 2:
|
|
await _apply_v2_migration(conn)
|
|
await conn.execute(
|
|
text(
|
|
"INSERT INTO bitable.bitable_meta (key, value, updated_at) "
|
|
"VALUES (:key, :value, NOW()) "
|
|
"ON CONFLICT (key) DO UPDATE SET value = :value, updated_at = NOW()"
|
|
),
|
|
{"key": _META_SCHEMA_VERSION_KEY, "value": str(_SCHEMA_VERSION)},
|
|
)
|
|
logger.info("bitable schema migrated: v%d → v%d", current_version, _SCHEMA_VERSION)
|
|
|
|
async def close(self) -> None:
|
|
"""Close engine, release all connections."""
|
|
if self._engine is not None:
|
|
await self._engine.dispose()
|
|
self._engine = None
|
|
self._session_factory = None
|
|
self._initialized = False
|
|
|
|
async def __aenter__(self) -> "BitableDB":
|
|
await self.init()
|
|
return self
|
|
|
|
async def __aexit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: TracebackType | None,
|
|
) -> None:
|
|
await self.close()
|
|
|
|
|
|
# Module-level singleton (initialized by app lifespan)
|
|
_db: BitableDB | None = None
|
|
|
|
|
|
async def init_bitable_db(database_url: str | None = None) -> BitableDB:
|
|
"""Initialize the bitable database (module-level singleton).
|
|
|
|
Called from ``app.py`` lifespan. On failure, raises — the caller should
|
|
catch and degrade gracefully (bitable API returns 503).
|
|
"""
|
|
global _db
|
|
if _db is not None and _db.is_initialized:
|
|
return _db
|
|
|
|
_db = BitableDB(database_url=database_url)
|
|
await _db.init()
|
|
logger.info("bitable DB initialized (schema version %d)", _SCHEMA_VERSION)
|
|
return _db
|
|
|
|
|
|
def get_bitable_db() -> BitableDB | None:
|
|
"""Return the module-level bitable DB singleton (or None if not init'd)."""
|
|
return _db
|
|
|
|
|
|
async def close_bitable_db() -> None:
|
|
"""Close the module-level bitable DB singleton."""
|
|
global _db
|
|
if _db is not None:
|
|
await _db.close()
|
|
_db = None
|