RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: make workspace_name nullable

This commit is contained in:
Rajat Ahuja 2026-01-14 13:18:03 -05:00
parent 8180c28c16
commit e5e2a3e082
8 changed files with 123 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

86
migrations/versions/119a52b73c60_support_external_embeddings.py
View File

@ -9,6 +9,8 @@ This migration:
message_embeddings tables for tracking vector store synchronization status.
4. Adds partial unique index on queue table for reconciler task deduplication,
ensuring only one pending reconciler task exists per work_unit_key.
5. Makes workspace_name nullable on queue table for system-level tasks (e.g., reconciler)
that don't belong to any specific workspace.
Revision ID: 119a52b73c60
Revises: 7c0d9a4e3b1f
@ -22,7 +24,7 @@ import sqlalchemy as sa
from alembic import op
from pgvector.sqlalchemy import Vector
from migrations.utils import column_exists, get_schema, index_exists
from migrations.utils import column_exists, constraint_exists, get_schema, index_exists
# revision identifiers, used by Alembic.
revision: str = "119a52b73c60"
@ -195,10 +197,92 @@ def upgrade() -> None:
postgresql_where=sa.text("task_type = 'reconciler' AND processed = false"),
)
# Make workspace_name nullable on queue table for system-level tasks
# This requires dropping and recreating the FK constraint
if constraint_exists("queue", "fk_queue_workspace_name", "foreignkey", inspector):
op.drop_constraint(
"fk_queue_workspace_name", "queue", type_="foreignkey", schema=schema
)
op.alter_column(
"queue",
"workspace_name",
existing_type=sa.TEXT(),
nullable=True,
schema=schema,
)
op.create_foreign_key(
"fk_queue_workspace_name",
"queue",
"workspaces",
["workspace_name"],
["name"],
source_schema=schema,
referent_schema=schema,
)
def downgrade() -> None:
"""Remove deleted_at columns and revert embedding columns."""
inspector = sa.inspect(op.get_bind())
conn = op.get_bind()
# Delete system-level queue items (with NULL workspace_name) before reverting to NOT NULL
# First delete any active_queue_sessions referencing these queue items
batch_size = 5000
# Delete active_queue_sessions for queue items with NULL workspace_name
while True:
result = conn.execute(
sa.text(
f"""
DELETE FROM "{schema}".active_queue_sessions
WHERE work_unit_key IN (
SELECT work_unit_key FROM "{schema}".queue
WHERE workspace_name IS NULL
)
LIMIT :batch_size
"""
),
{"batch_size": batch_size},
)
if result.rowcount == 0:
break
# Delete queue items with NULL workspace_name
while True:
result = conn.execute(
sa.text(
f"""
DELETE FROM "{schema}".queue
WHERE workspace_name IS NULL
LIMIT :batch_size
"""
),
{"batch_size": batch_size},
)
if result.rowcount == 0:
break
# Revert workspace_name to NOT NULL on queue table
op.drop_constraint(
"fk_queue_workspace_name", "queue", type_="foreignkey", schema=schema
)
op.alter_column(
"queue",
"workspace_name",
existing_type=sa.TEXT(),
nullable=False,
schema=schema,
)
op.create_foreign_key(
"fk_queue_workspace_name",
"queue",
"workspaces",
["workspace_name"],
["name"],
source_schema=schema,
referent_schema=schema,
)
# Drop reconciler queue index if it exists
if index_exists("queue", "uq_queue_work_unit_key", inspector):

32
src/deriver/consumer.py
View File

@ -34,6 +34,25 @@ async def process_item(queue_item: models.QueueItem) -> None:
queue_payload = queue_item.payload
workspace_name = queue_item.workspace_name
# Handle reconciler first - it's the only task type that doesn't require workspace_name
if task_type == "reconciler":
with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"):
try:
validated = ReconcilerPayload(**queue_payload)
except ValidationError as e:
logger.error(
"Invalid reconciler payload received: %s. Payload: %s",
str(e),
queue_payload,
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_reconciler(validated)
return
# All other task types require a workspace_name
if workspace_name is None:
raise ValueError(f"{task_type} tasks require a workspace_name")
if task_type == "webhook":
try:
validated = WebhookPayload(**queue_payload)
@ -122,19 +141,6 @@ async def process_item(queue_item: models.QueueItem) -> None:
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_deletion(validated, workspace_name)
elif task_type == "reconciler":
with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"):
try:
validated = ReconcilerPayload(**queue_payload)
except ValidationError as e:
logger.error(
"Invalid reconciler payload received: %s. Payload: %s",
str(e),
queue_payload,
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_reconciler(validated)
else:
raise ValueError(f"Invalid task type: {task_type}")

10
src/deriver/queue_manager.py
View File

@ -468,7 +468,10 @@ class QueueManager:
if removed and queue_item_count > 0:
# Only publish webhook if we actually removed an active session
try:
if work_unit.task_type in ["representation", "summary"]:
if (
work_unit.task_type in ["representation", "summary"]
and work_unit.workspace_name is not None
):
logger.debug(
f"Publishing queue.empty event for {work_unit_key} in workspace {work_unit.workspace_name}"
)
@ -730,7 +733,10 @@ class QueueManager:
)
await db.commit()
if work_unit.task_type in ["representation", "summary"]:
if (
work_unit.task_type in ["representation", "summary"]
and work_unit.workspace_name is not None
):
prometheus.DERIVER_QUEUE_ITEMS_PROCESSED.labels(
workspace_name=work_unit.workspace_name,
task_type=work_unit.task_type,

4
src/models.py
View File

@ -478,8 +478,8 @@ class QueueItem(Base):
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), index=True
)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), nullable=False, index=True
workspace_name: Mapped[str | None] = mapped_column(
ForeignKey("workspaces.name"), nullable=True, index=True
)
message_id: Mapped[int | None] = mapped_column(
BigInteger, ForeignKey("messages.id"), nullable=True

17
src/reconciler/scheduler.py
View File

@ -15,9 +15,7 @@ from datetime import datetime, timedelta, timezone
import sentry_sdk
from pydantic import BaseModel
from sqlalchemy import exists, select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.config import settings
@ -26,9 +24,6 @@ from src.models import QueueItem
logger = logging.getLogger(__name__)
# System workspace used for global reconciler tasks
SYSTEM_WORKSPACE_NAME = "__system__"
class ReconcilerTask(BaseModel):
"""Definition of a reconciler task."""
@ -216,9 +211,6 @@ class ReconcilerScheduler:
True if a task was enqueued, False if skipped
"""
async with tracked_db("reconciler_enqueue") as db:
# Ensure the system workspace exists (for FK constraint)
await self._ensure_system_workspace(db)
# Check if task is already in progress
in_progress_check = select(
exists(
@ -258,7 +250,7 @@ class ReconcilerScheduler:
},
session_id=None,
task_type="reconciler",
workspace_name=SYSTEM_WORKSPACE_NAME,
workspace_name=None,
message_id=None,
)
db.add(queue_item)
@ -274,10 +266,3 @@ class ReconcilerScheduler:
logger.info("Enqueued reconciler task: %s", task.name)
return True
async def _ensure_system_workspace(self, db: AsyncSession) -> None:
"""Ensure the system workspace exists for reconciler tasks."""
# Use upsert to create workspace if it doesn't exist
stmt = pg_insert(models.Workspace).values(name=SYSTEM_WORKSPACE_NAME)
stmt = stmt.on_conflict_do_nothing(index_elements=["name"])
await db.execute(stmt)

7
src/schemas.py
View File

@ -206,13 +206,6 @@ class WorkspaceCreate(WorkspaceBase):
model_config = ConfigDict(populate_by_name=True) # pyright: ignore
@field_validator("name")
@classmethod
def validate_not_reserved(cls, v: str) -> str:
if v == "__system__":
raise ValueError("Workspace name is reserved.")
return v
class WorkspaceGet(WorkspaceBase):
filters: dict[str, Any] | None = None

4
src/utils/work_unit.py
View File

@ -9,7 +9,7 @@ class ParsedWorkUnit(BaseModel):
"""Parsed work unit components."""
task_type: str
workspace_name: str
workspace_name: str | None
session_name: str | None
observer: str | None
observed: str | None
@ -149,7 +149,7 @@ def parse_work_unit_key(work_unit_key: str) -> ParsedWorkUnit:
)
return ParsedWorkUnit(
task_type=task_type,
workspace_name="__system__",
workspace_name=None,
session_name=None,
observer=None,
observed=None,

6
tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py
View File

@ -38,6 +38,9 @@ def prepare_support_external_embeddings(
verifier.assert_column_exists("message_embeddings", "last_sync_at", exists=False)
verifier.assert_column_exists("message_embeddings", "sync_attempts", exists=False)
# Queue workspace_name should be NOT NULL before migration
verifier.assert_column_exists("queue", "workspace_name", nullable=False)
# Indexes should not exist
verifier.assert_indexes_not_exist(INDEXES)
@ -64,5 +67,8 @@ def verify_support_external_embeddings(
verifier.assert_column_exists("message_embeddings", "last_sync_at", nullable=True)
verifier.assert_column_exists("message_embeddings", "sync_attempts", nullable=False)
# Queue workspace_name should now be nullable for system-level tasks
verifier.assert_column_exists("queue", "workspace_name", nullable=True)
# All indexes should exist
verifier.assert_indexes_exist(INDEXES)