Fix: fixed migration to handle postgres json properly

This commit is contained in:
weberm1 2025-06-05 13:58:30 +02:00
parent 749125c09a
commit 632478be3b

View File

@ -22,29 +22,63 @@
> migrator.drop_constraints(model, *constraints)
"""
"""Add rag_config field to knowledge table if not present."""
from contextlib import suppress
import peewee as pw
from peewee_migrate import Migrator
import peewee as pw
import json
with suppress(ImportError):
import playhouse.postgres_ext as pw_pext
class Knowledge(pw.Model):
class Meta:
table_name = 'knowledge'
# Try importing JSONField from playhouse.postgres_ext
with suppress(ImportError):
from playhouse.postgres_ext import JSONField as PostgresJSONField
# Fallback JSONField for SQLite (stores JSON as text)
class SQLiteJSONField(pw.TextField):
def db_value(self, value):
return json.dumps(value)
def python_value(self, value):
if value is not None:
return json.loads(value)
return None
def get_compatible_json_field(database: pw.Database):
"""Return a JSON-compatible field for the current database."""
if isinstance(database, pw.SqliteDatabase):
return SQLiteJSONField(null=False, default={"DEFAULT_RAG_SETTINGS": True})
else:
return PostgresJSONField(null=False, default={"DEFAULT_RAG_SETTINGS": True})
def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
"""Add rag_config column to knowledge table."""
Knowledge._meta.database = database # manually bind DB
"""Add rag_config JSON field to knowledge table"""
if 'knowledge' not in database.get_tables():
print("Knowledge table hasn't been created yet, skipping migration.")
return
class Knowledge(pw.Model):
class Meta:
table_name = 'knowledge'
Knowledge._meta.database = database # bind DB
migrator.add_fields(
Knowledge,
rag_config=pw.TextField(null=True, default=json.dumps({"DEFAULT_RAG_SETTINGS": True})),
rag_config=get_compatible_json_field(database)
)
def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
"""Remove rag_config column from knowledge table."""
Knowledge._meta.database = database # manually bind DB
"""Remove rag_config field from knowledge table."""
if 'knowledge' not in database.get_tables():
print("Knowledge table hasn't been created yet, skipping migration.")
return
migrator.remove_fields(Knowledge, "rag_config")
class Knowledge(pw.Model):
class Meta:
table_name = 'knowledge'
Knowledge._meta.database = database
migrator.remove_fields(Knowledge, 'rag_config')