PowerOn/gateway
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore: fix pydantic v2 issue

Browse source
This commit is contained in:
Christopher Gondek 2025-10-03 16:51:43 +02:00
parent 4bfeded9d0
commit 37f01a2156
1 changed files with 362 additions and 227 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

329
modules/connectors/connectorDbPostgre.py
View file

@ -18,60 +18,100 @@ logger = logging.getLogger(__name__)
# No mapping needed - table name = Pydantic model name exactly # No mapping needed - table name = Pydantic model name exactly
class SystemTable(BaseModel, ModelMixin): class SystemTable(BaseModel, ModelMixin):
"""Data model for system table entries""" """Data model for system table entries"""
table_name: str = Field( table_name: str = Field(
description="Name of the table", description="Name of the table",
frontend_type="text", frontend_type="text",
frontend_readonly=True, frontend_readonly=True,
frontend_required=True frontend_required=True,
) )
initial_id: Optional[str] = Field( initial_id: Optional[str] = Field(
default=None, default=None,
description="Initial ID for the table", description="Initial ID for the table",
frontend_type="text", frontend_type="text",
frontend_readonly=True, frontend_readonly=True,
frontend_required=False frontend_required=False,
) )
def _get_model_fields(model_class) -> Dict[str, str]: def _get_model_fields(model_class) -> Dict[str, str]:
"""Get all fields from Pydantic model and map to SQL types.""" """Get all fields from Pydantic model and map to SQL types."""
if not hasattr(model_class, '__fields__'): # Pydantic v2 uses model_fields instead of __fields__
if hasattr(model_class, "model_fields"):
model_fields = model_class.model_fields
elif hasattr(model_class, "__fields__"):
model_fields = model_class.__fields__
else:
return {} return {}
fields = {} fields = {}
for field_name, field_info in model_class.__fields__.items(): for field_name, field_info in model_fields.items():
field_type = field_info.type_ # Pydantic v2 uses annotation instead of type_
field_type = (
field_info.annotation
if hasattr(field_info, "annotation")
else field_info.type_
)
# Check for JSONB fields (Dict, List, or complex types) # Check for JSONB fields (Dict, List, or complex types)
if (field_type == dict or if (
field_type == list or field_type == dict
(hasattr(field_type, '__origin__') and field_type.__origin__ in (dict, list)) or or field_type == list
field_name in ['execParameters', 'expectedDocumentFormats', 'resultDocuments', 'logs', 'messages', 'stats', 'tasks']): or (
fields[field_name] = 'JSONB' hasattr(field_type, "__origin__")
and field_type.__origin__ in (dict, list)
)
or field_name
in [
"execParameters",
"expectedDocumentFormats",
"resultDocuments",
"logs",
"messages",
"stats",
"tasks",
]
):
fields[field_name] = "JSONB"
# Simple type mapping # Simple type mapping
elif field_type in (str, type(None)) or (get_origin(field_type) is Union and type(None) in get_args(field_type)): elif field_type in (str, type(None)) or (
fields[field_name] = 'TEXT' get_origin(field_type) is Union and type(None) in get_args(field_type)
):
fields[field_name] = "TEXT"
elif field_type == int: elif field_type == int:
fields[field_name] = 'INTEGER' fields[field_name] = "INTEGER"
elif field_type == float: elif field_type == float:
fields[field_name] = 'DOUBLE PRECISION' fields[field_name] = "DOUBLE PRECISION"
elif field_type == bool: elif field_type == bool:
fields[field_name] = 'BOOLEAN' fields[field_name] = "BOOLEAN"
else: else:
fields[field_name] = 'TEXT' # Default to TEXT fields[field_name] = "TEXT" # Default to TEXT
return fields return fields
# No caching needed with proper database # No caching needed with proper database
class DatabaseConnector: class DatabaseConnector:
""" """
A connector for PostgreSQL-based data storage. A connector for PostgreSQL-based data storage.
Provides generic database operations without user/mandate filtering. Provides generic database operations without user/mandate filtering.
Uses PostgreSQL with JSONB columns for flexible data storage. Uses PostgreSQL with JSONB columns for flexible data storage.
""" """
def __init__(self, dbHost: str, dbDatabase: str, dbUser: str = None, dbPassword: str = None, dbPort: int = None, userId: str = None):
def __init__(
self,
dbHost: str,
dbDatabase: str,
dbUser: str = None,
dbPassword: str = None,
dbPort: int = None,
userId: str = None,
):
# Store the input parameters # Store the input parameters
self.dbHost = dbHost self.dbHost = dbHost
self.dbDatabase = dbDatabase self.dbDatabase = dbDatabase
@ -95,7 +135,6 @@ class DatabaseConnector:
self._systemTableName = "_system" self._systemTableName = "_system"
self._initializeSystemTable() self._initializeSystemTable()
def initDbSystem(self): def initDbSystem(self):
"""Initialize the database system - creates database and tables.""" """Initialize the database system - creates database and tables."""
try: try:
@ -123,13 +162,15 @@ class DatabaseConnector:
database="postgres", database="postgres",
user=self.dbUser, user=self.dbUser,
password=self.dbPassword, password=self.dbPassword,
client_encoding='utf8' client_encoding="utf8",
) )
conn.autocommit = True conn.autocommit = True
with conn.cursor() as cursor: with conn.cursor() as cursor:
# Check if database exists # Check if database exists
cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (self.dbDatabase,)) cursor.execute(
"SELECT 1 FROM pg_database WHERE datname = %s", (self.dbDatabase,)
)
exists = cursor.fetchone() exists = cursor.fetchone()
if not exists: if not exists:
@ -143,8 +184,9 @@ class DatabaseConnector:
except Exception as e: except Exception as e:
logger.error(f"FATAL ERROR: Cannot create database: {e}") logger.error(f"FATAL ERROR: Cannot create database: {e}")
logger.error("Database connection failed - application cannot start") logger.error("Database connection failed - application cannot start")
raise RuntimeError(f"FATAL ERROR: Cannot create database '{self.dbDatabase}': {e}") raise RuntimeError(
f"FATAL ERROR: Cannot create database '{self.dbDatabase}': {e}"
)
def _create_tables(self): def _create_tables(self):
"""Create only the system table - application tables are created by interfaces.""" """Create only the system table - application tables are created by interfaces."""
@ -156,7 +198,7 @@ class DatabaseConnector:
database=self.dbDatabase, database=self.dbDatabase,
user=self.dbUser, user=self.dbUser,
password=self.dbPassword, password=self.dbPassword,
client_encoding='utf8' client_encoding="utf8",
) )
conn.autocommit = True conn.autocommit = True
@ -175,7 +217,9 @@ class DatabaseConnector:
except Exception as e: except Exception as e:
logger.error(f"FATAL ERROR: Cannot create system table: {e}") logger.error(f"FATAL ERROR: Cannot create system table: {e}")
logger.error("Database system table creation failed - application cannot start") logger.error(
"Database system table creation failed - application cannot start"
)
raise RuntimeError(f"FATAL ERROR: Cannot create system table: {e}") raise RuntimeError(f"FATAL ERROR: Cannot create system table: {e}")
def _connect(self): def _connect(self):
@ -188,8 +232,8 @@ class DatabaseConnector:
database=self.dbDatabase, database=self.dbDatabase,
user=self.dbUser, user=self.dbUser,
password=self.dbPassword, password=self.dbPassword,
client_encoding='utf8', client_encoding="utf8",
cursor_factory=psycopg2.extras.RealDictCursor cursor_factory=psycopg2.extras.RealDictCursor,
) )
self.connection.autocommit = False # Use transactions self.connection.autocommit = False # Use transactions
except Exception as e: except Exception as e:
@ -219,7 +263,7 @@ class DatabaseConnector:
# Check if system table has any data # Check if system table has any data
cursor.execute('SELECT COUNT(*) FROM "_system"') cursor.execute('SELECT COUNT(*) FROM "_system"')
row = cursor.fetchone() row = cursor.fetchone()
count = row['count'] if row else 0 count = row["count"] if row else 0
self.connection.commit() self.connection.commit()
except Exception as e: except Exception as e:
@ -236,7 +280,7 @@ class DatabaseConnector:
system_data = {} system_data = {}
for row in rows: for row in rows:
system_data[row['table_name']] = row['initial_id'] system_data[row["table_name"]] = row["initial_id"]
return system_data return system_data
except Exception as e: except Exception as e:
@ -252,10 +296,13 @@ class DatabaseConnector:
# Insert new data # Insert new data
for table_name, initial_id in data.items(): for table_name, initial_id in data.items():
cursor.execute(""" cursor.execute(
"""
INSERT INTO "_system" ("table_name", "initial_id", "_modifiedAt") INSERT INTO "_system" ("table_name", "initial_id", "_modifiedAt")
VALUES (%s, %s, %s) VALUES (%s, %s, %s)
""", (table_name, initial_id, get_utc_timestamp())) """,
(table_name, initial_id, get_utc_timestamp()),
)
self.connection.commit() self.connection.commit()
return True return True
@ -271,8 +318,11 @@ class DatabaseConnector:
with self.connection.cursor() as cursor: with self.connection.cursor() as cursor:
# Check if system table exists # Check if system table exists
cursor.execute("SELECT COUNT(*) FROM pg_stat_user_tables WHERE relname = %s", (self._systemTableName,)) cursor.execute(
exists = cursor.fetchone()['count'] > 0 "SELECT COUNT(*) FROM pg_stat_user_tables WHERE relname = %s",
(self._systemTableName,),
)
exists = cursor.fetchone()["count"] > 0
if not exists: if not exists:
# Create system table # Create system table
@ -287,14 +337,19 @@ class DatabaseConnector:
logger.info("System table created successfully") logger.info("System table created successfully")
else: else:
# Check if we need to add missing columns to existing table # Check if we need to add missing columns to existing table
cursor.execute(""" cursor.execute(
"""
SELECT column_name FROM information_schema.columns SELECT column_name FROM information_schema.columns
WHERE table_name = %s AND table_schema = 'public' WHERE table_name = %s AND table_schema = 'public'
""", (self._systemTableName,)) """,
existing_columns = [row['column_name'] for row in cursor.fetchall()] (self._systemTableName,),
)
existing_columns = [row["column_name"] for row in cursor.fetchall()]
if '_modifiedAt' not in existing_columns: if "_modifiedAt" not in existing_columns:
cursor.execute(f'ALTER TABLE "{self._systemTableName}" ADD COLUMN "_modifiedAt" DOUBLE PRECISION') cursor.execute(
f'ALTER TABLE "{self._systemTableName}" ADD COLUMN "_modifiedAt" DOUBLE PRECISION'
)
return True return True
except Exception as e: except Exception as e:
@ -314,26 +369,30 @@ class DatabaseConnector:
with self.connection.cursor() as cursor: with self.connection.cursor() as cursor:
# Check if table exists by querying information_schema with case-insensitive search # Check if table exists by querying information_schema with case-insensitive search
cursor.execute(''' cursor.execute(
"""
SELECT COUNT(*) FROM information_schema.tables SELECT COUNT(*) FROM information_schema.tables
WHERE LOWER(table_name) = LOWER(%s) AND table_schema = 'public' WHERE LOWER(table_name) = LOWER(%s) AND table_schema = 'public'
''', (table,)) """,
exists = cursor.fetchone()['count'] > 0 (table,),
)
exists = cursor.fetchone()["count"] > 0
if not exists: if not exists:
# Create table from Pydantic model # Create table from Pydantic model
self._create_table_from_model(cursor, table, model_class) self._create_table_from_model(cursor, table, model_class)
logger.info(f"Created table '{table}' with columns from Pydantic model") logger.info(
f"Created table '{table}' with columns from Pydantic model"
)
self.connection.commit() self.connection.commit()
return True return True
except Exception as e: except Exception as e:
logger.error(f"Error ensuring table {table} exists: {e}") logger.error(f"Error ensuring table {table} exists: {e}")
if hasattr(self, 'connection') and self.connection: if hasattr(self, "connection") and self.connection:
self.connection.rollback() self.connection.rollback()
return False return False
def _create_table_from_model(self, cursor, table: str, model_class: type) -> None: def _create_table_from_model(self, cursor, table: str, model_class: type) -> None:
"""Create table with columns matching Pydantic model fields.""" """Create table with columns matching Pydantic model fields."""
fields = _get_model_fields(model_class) fields = _get_model_fields(model_class)
@ -341,16 +400,18 @@ class DatabaseConnector:
# Build column definitions with quoted identifiers to preserve exact case # Build column definitions with quoted identifiers to preserve exact case
columns = ['"id" VARCHAR(255) PRIMARY KEY'] columns = ['"id" VARCHAR(255) PRIMARY KEY']
for field_name, sql_type in fields.items(): for field_name, sql_type in fields.items():
if field_name != 'id': # Skip id, already defined if field_name != "id": # Skip id, already defined
columns.append(f'"{field_name}" {sql_type}') columns.append(f'"{field_name}" {sql_type}')
# Add metadata columns # Add metadata columns
columns.extend([ columns.extend(
[
'"_createdAt" DOUBLE PRECISION', '"_createdAt" DOUBLE PRECISION',
'"_modifiedAt" DOUBLE PRECISION', '"_modifiedAt" DOUBLE PRECISION',
'"_createdBy" VARCHAR(255)', '"_createdBy" VARCHAR(255)',
'"_modifiedBy" VARCHAR(255)' '"_modifiedBy" VARCHAR(255)',
]) ]
)
# Create table # Create table
sql = f'CREATE TABLE IF NOT EXISTS "{table}" ({", ".join(columns)})' sql = f'CREATE TABLE IF NOT EXISTS "{table}" ({", ".join(columns)})'
@ -358,16 +419,27 @@ class DatabaseConnector:
# Create indexes for foreign keys # Create indexes for foreign keys
for field_name in fields: for field_name in fields:
if field_name.endswith('Id') and field_name != 'id': if field_name.endswith("Id") and field_name != "id":
cursor.execute(f'CREATE INDEX IF NOT EXISTS "idx_{table}_{field_name}" ON "{table}" ("{field_name}")') cursor.execute(
f'CREATE INDEX IF NOT EXISTS "idx_{table}_{field_name}" ON "{table}" ("{field_name}")'
)
def _save_record(
def _save_record(self, cursor, table: str, recordId: str, record: Dict[str, Any], model_class: type) -> None: self,
cursor,
table: str,
recordId: str,
record: Dict[str, Any],
model_class: type,
) -> None:
"""Save record to normalized table with explicit columns.""" """Save record to normalized table with explicit columns."""
# Get columns from Pydantic model instead of database schema # Get columns from Pydantic model instead of database schema
fields = _get_model_fields(model_class) fields = _get_model_fields(model_class)
columns = ['id'] + [field for field in fields.keys() if field != 'id'] + ['_createdAt', '_createdBy', '_modifiedAt', '_modifiedBy'] columns = (
["id"]
+ [field for field in fields.keys() if field != "id"]
+ ["_createdAt", "_createdBy", "_modifiedAt", "_modifiedBy"]
)
if not columns: if not columns:
logger.error(f"No columns found for table {table}") logger.error(f"No columns found for table {table}")
@ -377,7 +449,7 @@ class DatabaseConnector:
filtered_record = {k: v for k, v in record.items() if k in columns} filtered_record = {k: v for k, v in record.items() if k in columns}
# Ensure id is set # Ensure id is set
filtered_record['id'] = recordId filtered_record["id"] = recordId
# Prepare values in the correct order # Prepare values in the correct order
values = [] values = []
@ -385,7 +457,7 @@ class DatabaseConnector:
value = filtered_record.get(col) value = filtered_record.get(col)
# Handle timestamp fields - store as Unix timestamps (floats) for consistency # Handle timestamp fields - store as Unix timestamps (floats) for consistency
if col in ['_createdAt', '_modifiedAt'] and value is not None: if col in ["_createdAt", "_modifiedAt"] and value is not None:
if isinstance(value, str): if isinstance(value, str):
# Try to parse string as timestamp # Try to parse string as timestamp
try: try:
@ -394,12 +466,13 @@ class DatabaseConnector:
pass # Keep as string if parsing fails pass # Keep as string if parsing fails
# Convert enum values to their string representation # Convert enum values to their string representation
elif hasattr(value, 'value'): elif hasattr(value, "value"):
value = value.value value = value.value
# Handle JSONB fields - ensure proper JSON format for PostgreSQL # Handle JSONB fields - ensure proper JSON format for PostgreSQL
elif col in fields and fields[col] == 'JSONB' and value is not None: elif col in fields and fields[col] == "JSONB" and value is not None:
import json import json
if isinstance(value, (dict, list)): if isinstance(value, (dict, list)):
# Convert Python objects to JSON string for PostgreSQL JSONB # Convert Python objects to JSON string for PostgreSQL JSONB
value = json.dumps(value) value = json.dumps(value)
@ -419,11 +492,16 @@ class DatabaseConnector:
values.append(value) values.append(value)
# Build INSERT/UPDATE with quoted identifiers # Build INSERT/UPDATE with quoted identifiers
col_names = ', '.join([f'"{col}"' for col in columns]) col_names = ", ".join([f'"{col}"' for col in columns])
placeholders = ', '.join(['%s'] * len(columns)) placeholders = ", ".join(["%s"] * len(columns))
updates = ', '.join([f'"{col}" = EXCLUDED."{col}"' for col in columns[1:] if col not in ['_createdAt', '_createdBy']]) updates = ", ".join(
[
f'"{col}" = EXCLUDED."{col}"'
for col in columns[1:]
if col not in ["_createdAt", "_createdBy"]
]
)
sql = f'INSERT INTO "{table}" ({col_names}) VALUES ({placeholders}) ON CONFLICT ("id") DO UPDATE SET {updates}' sql = f'INSERT INTO "{table}" ({col_names}) VALUES ({placeholders}) ON CONFLICT ("id") DO UPDATE SET {updates}'
@ -447,11 +525,15 @@ class DatabaseConnector:
record = dict(row) record = dict(row)
fields = _get_model_fields(model_class) fields = _get_model_fields(model_class)
# Parse JSONB fields back to Python objects # Parse JSONB fields back to Python objects
for field_name, field_type in fields.items(): for field_name, field_type in fields.items():
if field_type == 'JSONB' and field_name in record and record[field_name] is not None: if (
field_type == "JSONB"
and field_name in record
and record[field_name] is not None
):
import json import json
try: try:
if isinstance(record[field_name], str): if isinstance(record[field_name], str):
# Parse JSON string back to Python object # Parse JSON string back to Python object
@ -464,7 +546,9 @@ class DatabaseConnector:
record[field_name] = json.loads(str(record[field_name])) record[field_name] = json.loads(str(record[field_name]))
except (json.JSONDecodeError, TypeError, ValueError): except (json.JSONDecodeError, TypeError, ValueError):
# If parsing fails, keep as string # If parsing fails, keep as string
logger.warning(f"Could not parse JSONB field {field_name}, keeping as string: {record[field_name]}") logger.warning(
f"Could not parse JSONB field {field_name}, keeping as string: {record[field_name]}"
)
pass pass
return record return record
@ -472,7 +556,9 @@ class DatabaseConnector:
logger.error(f"Error loading record {recordId} from table {table}: {e}") logger.error(f"Error loading record {recordId} from table {table}: {e}")
return None return None
def _saveRecord(self, model_class: type, recordId: str, record: Dict[str, Any]) -> bool: def _saveRecord(
self, model_class: type, recordId: str, record: Dict[str, Any]
) -> bool:
"""Saves a single record to the table.""" """Saves a single record to the table."""
table = model_class.__name__ table = model_class.__name__
@ -521,30 +607,43 @@ class DatabaseConnector:
fields = _get_model_fields(model_class) fields = _get_model_fields(model_class)
for record in records: for record in records:
for field_name, field_type in fields.items(): for field_name, field_type in fields.items():
if field_type == 'JSONB' and field_name in record: if field_type == "JSONB" and field_name in record:
if record[field_name] is None: if record[field_name] is None:
# Convert None to appropriate default based on field name # Convert None to appropriate default based on field name
if field_name in ['logs', 'messages', 'tasks', 'expectedDocumentFormats', 'resultDocuments']: if field_name in [
"logs",
"messages",
"tasks",
"expectedDocumentFormats",
"resultDocuments",
]:
record[field_name] = [] record[field_name] = []
elif field_name in ['execParameters', 'stats']: elif field_name in ["execParameters", "stats"]:
record[field_name] = {} record[field_name] = {}
else: else:
record[field_name] = None record[field_name] = None
else: else:
import json import json
try: try:
if isinstance(record[field_name], str): if isinstance(record[field_name], str):
# Parse JSON string back to Python object # Parse JSON string back to Python object
record[field_name] = json.loads(record[field_name]) record[field_name] = json.loads(
record[field_name]
)
elif isinstance(record[field_name], (dict, list)): elif isinstance(record[field_name], (dict, list)):
# Already a Python object, keep as is # Already a Python object, keep as is
pass pass
else: else:
# Try to parse as JSON # Try to parse as JSON
record[field_name] = json.loads(str(record[field_name])) record[field_name] = json.loads(
str(record[field_name])
)
except (json.JSONDecodeError, TypeError, ValueError): except (json.JSONDecodeError, TypeError, ValueError):
# If parsing fails, keep as string # If parsing fails, keep as string
logger.warning(f"Could not parse JSONB field {field_name}, keeping as string: {record[field_name]}") logger.warning(
f"Could not parse JSONB field {field_name}, keeping as string: {record[field_name]}"
)
pass pass
return records return records
@ -552,8 +651,6 @@ class DatabaseConnector:
logger.error(f"Error loading table {table}: {e}") logger.error(f"Error loading table {table}: {e}")
return [] return []
def _registerInitialId(self, table: str, initialId: str) -> bool: def _registerInitialId(self, table: str, initialId: str) -> bool:
"""Registers the initial ID for a table.""" """Registers the initial ID for a table."""
try: try:
@ -568,13 +665,17 @@ class DatabaseConnector:
else: else:
# Check if the existing initial ID still exists in the table # Check if the existing initial ID still exists in the table
existingInitialId = systemData[table] existingInitialId = systemData[table]
records = self.getRecordset(model_class, recordFilter={"id": existingInitialId}) records = self.getRecordset(
model_class, recordFilter={"id": existingInitialId}
)
if not records: if not records:
# The initial record no longer exists, update to the new one # The initial record no longer exists, update to the new one
systemData[table] = initialId systemData[table] = initialId
success = self._saveSystemTable(systemData) success = self._saveSystemTable(systemData)
if success: if success:
logger.info(f"Initial ID updated from {existingInitialId} to {initialId} for table {table}") logger.info(
f"Initial ID updated from {existingInitialId} to {initialId} for table {table}"
)
return success return success
else: else:
return True return True
@ -591,7 +692,9 @@ class DatabaseConnector:
del systemData[table] del systemData[table]
success = self._saveSystemTable(systemData) success = self._saveSystemTable(systemData)
if success: if success:
logger.info(f"Initial ID for table {table} removed from system table") logger.info(
f"Initial ID for table {table} removed from system table"
)
return success return success
return True # If not present, this is not an error return True # If not present, this is not an error
except Exception as e: except Exception as e:
@ -628,7 +731,7 @@ class DatabaseConnector:
ORDER BY table_name ORDER BY table_name
""") """)
rows = cursor.fetchall() rows = cursor.fetchall()
tables = [row['table_name'] for row in rows] tables = [row["table_name"] for row in rows]
except Exception as e: except Exception as e:
logger.error(f"Error reading the database {self.dbDatabase}: {e}") logger.error(f"Error reading the database {self.dbDatabase}: {e}")
@ -645,7 +748,9 @@ class DatabaseConnector:
return fields return fields
def getSchema(self, model_class: type, language: str = None) -> Dict[str, Dict[str, Any]]: def getSchema(
self, model_class: type, language: str = None
) -> Dict[str, Dict[str, Any]]:
"""Returns a schema object for a table with data types and labels.""" """Returns a schema object for a table with data types and labels."""
data = self._loadTable(model_class) data = self._loadTable(model_class)
@ -660,14 +765,16 @@ class DatabaseConnector:
dataType = type(value).__name__ dataType = type(value).__name__
label = field label = field
schema[field] = { schema[field] = {"type": dataType, "label": label}
"type": dataType,
"label": label
}
return schema return schema
def getRecordset(self, model_class: type, fieldFilter: List[str] = None, recordFilter: Dict[str, Any] = None) -> List[Dict[str, Any]]: def getRecordset(
self,
model_class: type,
fieldFilter: List[str] = None,
recordFilter: Dict[str, Any] = None,
) -> List[Dict[str, Any]]:
"""Returns a list of records from a table, filtered by criteria.""" """Returns a list of records from a table, filtered by criteria."""
table = model_class.__name__ table = model_class.__name__
@ -700,30 +807,43 @@ class DatabaseConnector:
fields = _get_model_fields(model_class) fields = _get_model_fields(model_class)
for record in records: for record in records:
for field_name, field_type in fields.items(): for field_name, field_type in fields.items():
if field_type == 'JSONB' and field_name in record: if field_type == "JSONB" and field_name in record:
if record[field_name] is None: if record[field_name] is None:
# Convert None to appropriate default based on field name # Convert None to appropriate default based on field name
if field_name in ['logs', 'messages', 'tasks', 'expectedDocumentFormats', 'resultDocuments']: if field_name in [
"logs",
"messages",
"tasks",
"expectedDocumentFormats",
"resultDocuments",
]:
record[field_name] = [] record[field_name] = []
elif field_name in ['execParameters', 'stats']: elif field_name in ["execParameters", "stats"]:
record[field_name] = {} record[field_name] = {}
else: else:
record[field_name] = None record[field_name] = None
else: else:
import json import json
try: try:
if isinstance(record[field_name], str): if isinstance(record[field_name], str):
# Parse JSON string back to Python object # Parse JSON string back to Python object
record[field_name] = json.loads(record[field_name]) record[field_name] = json.loads(
record[field_name]
)
elif isinstance(record[field_name], (dict, list)): elif isinstance(record[field_name], (dict, list)):
# Already a Python object, keep as is # Already a Python object, keep as is
pass pass
else: else:
# Try to parse as JSON # Try to parse as JSON
record[field_name] = json.loads(str(record[field_name])) record[field_name] = json.loads(
str(record[field_name])
)
except (json.JSONDecodeError, TypeError, ValueError): except (json.JSONDecodeError, TypeError, ValueError):
# If parsing fails, keep as string # If parsing fails, keep as string
logger.warning(f"Could not parse JSONB field {field_name}, keeping as string: {record[field_name]}") logger.warning(
f"Could not parse JSONB field {field_name}, keeping as string: {record[field_name]}"
)
pass pass
# If fieldFilter is available, reduce the fields # If fieldFilter is available, reduce the fields
@ -742,7 +862,9 @@ class DatabaseConnector:
logger.error(f"Error loading records from table {table}: {e}") logger.error(f"Error loading records from table {table}: {e}")
return [] return []
def recordCreate(self, model_class: type, record: Union[Dict[str, Any], BaseModel]) -> Dict[str, Any]: def recordCreate(
self, model_class: type, record: Union[Dict[str, Any], BaseModel]
) -> Dict[str, Any]:
"""Creates a new record in a table based on Pydantic model class.""" """Creates a new record in a table based on Pydantic model class."""
# If record is a Pydantic model, convert to dict # If record is a Pydantic model, convert to dict
if isinstance(record, BaseModel): if isinstance(record, BaseModel):
@ -769,7 +891,9 @@ class DatabaseConnector:
return record return record
def recordModify(self, model_class: type, recordId: str, record: Union[Dict[str, Any], BaseModel]) -> Dict[str, Any]: def recordModify(
self, model_class: type, recordId: str, record: Union[Dict[str, Any], BaseModel]
) -> Dict[str, Any]:
"""Modifies an existing record in a table based on Pydantic model class.""" """Modifies an existing record in a table based on Pydantic model class."""
# Load existing record # Load existing record
existingRecord = self._loadRecord(model_class, recordId) existingRecord = self._loadRecord(model_class, recordId)
@ -787,8 +911,12 @@ class DatabaseConnector:
# CRITICAL: Ensure we never modify the ID # CRITICAL: Ensure we never modify the ID
if "id" in record and str(record["id"]) != recordId: if "id" in record and str(record["id"]) != recordId:
logger.error(f"Attempted to modify record ID from {recordId} to {record['id']}") logger.error(
raise ValueError("Cannot modify record ID - it must match the provided recordId") f"Attempted to modify record ID from {recordId} to {record['id']}"
)
raise ValueError(
"Cannot modify record ID - it must match the provided recordId"
)
# Update existing record with new data # Update existing record with new data
existingRecord.update(record) existingRecord.update(record)
@ -807,7 +935,9 @@ class DatabaseConnector:
with self.connection.cursor() as cursor: with self.connection.cursor() as cursor:
# Check if record exists # Check if record exists
cursor.execute(f'SELECT "id" FROM "{table}" WHERE "id" = %s', (recordId,)) cursor.execute(
f'SELECT "id" FROM "{table}" WHERE "id" = %s', (recordId,)
)
if not cursor.fetchone(): if not cursor.fetchone():
return False return False
@ -815,7 +945,9 @@ class DatabaseConnector:
initialId = self.getInitialId(model_class) initialId = self.getInitialId(model_class)
if initialId is not None and initialId == recordId: if initialId is not None and initialId == recordId:
self._removeInitialId(table) self._removeInitialId(table)
logger.info(f"Initial ID {recordId} for table {table} has been removed from the system table") logger.info(
f"Initial ID {recordId} for table {table} has been removed from the system table"
)
# Delete the record # Delete the record
cursor.execute(f'DELETE FROM "{table}" WHERE "id" = %s', (recordId,)) cursor.execute(f'DELETE FROM "{table}" WHERE "id" = %s', (recordId,))
@ -830,7 +962,6 @@ class DatabaseConnector:
self.connection.rollback() self.connection.rollback()
return False return False
def getInitialId(self, model_class: type) -> Optional[str]: def getInitialId(self, model_class: type) -> Optional[str]:
"""Returns the initial ID for a table.""" """Returns the initial ID for a table."""
table = model_class.__name__ table = model_class.__name__
@ -840,7 +971,11 @@ class DatabaseConnector:
def close(self): def close(self):
"""Close the database connection.""" """Close the database connection."""
if hasattr(self, 'connection') and self.connection and not self.connection.closed: if (
hasattr(self, "connection")
and self.connection
and not self.connection.closed
):
self.connection.close() self.connection.close()
def __del__(self): def __del__(self):