66from sqlalchemy .orm import sessionmaker , declarative_base , relationship
77from sqlalchemy .exc import SQLAlchemyError
88from sqlalchemy .pool import QueuePool
9- from sqlalchemy .dialects .postgresql import JSONB
9+ from sqlalchemy .types import TypeDecorator , JSON
10+ from sqlalchemy .ext .mutable import MutableDict
1011from contextlib import contextmanager
1112
13+ # Import JSONB specifically for PostgreSQL dialect
14+ try :
15+ from sqlalchemy .dialects .postgresql import JSONB
16+ except ImportError :
17+ # Fallback for environments where psycopg2/psycopg is not installed
18+ # or when running against non-PostgreSQL databases like SQLite
19+ JSONB = None
20+
1221Base = declarative_base ()
1322_logger = logging .getLogger (__name__ )
1423
24+ class JSONBOrJSON (TypeDecorator ):
25+ """
26+ A TypeDecorator that uses JSONB for PostgreSQL and JSON for other databases.
27+ This provides cross-database compatibility for JSON column types.
28+ """
29+ impl = JSON # Default implementation for non-PostgreSQL dialects
30+
31+ cache_ok = True # Indicate that this type is safe to cache
32+
33+ def load_dialect_impl (self , dialect ):
34+ if dialect .name == 'postgresql' and JSONB is not None :
35+ return dialect .type_descriptor (JSONB ())
36+ else :
37+ return dialect .type_descriptor (JSON ())
38+
39+ def process_bind_param (self , value , dialect ):
40+ # No special processing needed for binding, SQLAlchemy handles JSON serialization
41+ return value
42+
43+ def process_result_value (self , value , dialect ):
44+ # No special processing needed for results, SQLAlchemy handles JSON deserialization
45+ return value
46+
47+ # To make the JSON column mutable (i.e., changes to the dictionary are detected)
48+ MutableJSONBOrJSON = MutableDict .as_mutable (JSONBOrJSON )
49+
50+
1551class Project (Base ):
1652 __tablename__ = 'projects'
1753 id = Column (String , primary_key = True , index = True )
@@ -29,7 +65,8 @@ class Pattern(Base):
2965 id = Column (String , primary_key = True , index = True )
3066 project_id = Column (String , ForeignKey ('projects.id' ), nullable = True ) # Optional link to project
3167 document_text = Column (Text , nullable = False )
32- metadata_json = Column (JSONB , nullable = False , default = {})
68+ # Use MutableJSONBOrJSON for cross-database compatibility and mutability
69+ metadata_json = Column (MutableJSONBOrJSON , nullable = False , default = {})
3370 created_at = Column (DateTime , default = datetime .utcnow )
3471 updated_at = Column (DateTime , default = datetime .utcnow , onupdate = datetime .utcnow )
3572
@@ -46,7 +83,8 @@ class ErrorSolution(Base):
4683 id = Column (String , primary_key = True , index = True )
4784 project_id = Column (String , ForeignKey ('projects.id' ), nullable = True ) # Optional link to project
4885 document_text = Column (Text , nullable = False )
49- metadata_json = Column (JSONB , nullable = False , default = {})
86+ # Use MutableJSONBOrJSON for cross-database compatibility and mutability
87+ metadata_json = Column (MutableJSONBOrJSON , nullable = False , default = {})
5088 created_at = Column (DateTime , default = datetime .utcnow )
5189 updated_at = Column (DateTime , default = datetime .utcnow , onupdate = datetime .utcnow )
5290
@@ -258,15 +296,15 @@ def filter_records(self, model: Base, filters: Dict[str, Any], limit: Optional[i
258296 return []
259297
260298 def search_records_by_metadata (self , model : Base , metadata_filter : Dict [str , Any ], limit : Optional [int ] = None ) -> List [Dict [str , Any ]]:
261- """Search records by JSONB metadata fields."""
299+ """Search records by JSONB/JSON metadata fields using cross-database compatible operators ."""
262300 try :
263301 with self .get_db_session () as session :
264302 query = session .query (model )
265303
266- # Apply metadata filters using JSONB operators
304+ # Apply metadata filters using the cross-database compatible .contains() operator
267305 for key , value in metadata_filter .items ():
268- # Use JSONB contains operator for nested key-value searches
269- filter_condition = model .metadata_json .op ( '@>' ) ({key : value })
306+ # . contains() works for both PostgreSQL JSONB and SQLite JSON
307+ filter_condition = model .metadata_json .contains ({key : value })
270308 query = query .filter (filter_condition )
271309
272310 if limit :
@@ -354,4 +392,3 @@ def reset_db(self) -> bool:
354392 self ._logger .error (f"Failed to reset PostgreSQL database: { e } " )
355393 return False
356394 return False
357-
0 commit comments