109 lines
4 KiB
Python
109 lines
4 KiB
Python
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy import create_engine
|
|
from typing import Type, Dict, Any
|
|
from sqlalchemy.exc import IntegrityError, SQLAlchemyError
|
|
from sqlalchemy.orm import Session
|
|
|
|
from core.errors.logger import logger
|
|
from core.models.manage.session_management import get_session
|
|
from core.models.orm_models.Base import BaseModel
|
|
|
|
from core.models.manage.wrapper import safe_db_operation, WrapperRollback
|
|
from core.models.DatabaseOperation import DatabaseOperation, DBErrorType
|
|
|
|
# all models it knows how to do:
|
|
from core.models.orm_models.CachedSync import CachedSync
|
|
from core.models.orm_models.Location import Location
|
|
from core.models.orm_models.Operator import Operator
|
|
from core.models.orm_models.Dependency import Dependency
|
|
|
|
from core.models.orm_models.EncryptedProxy import EncryptedProxy
|
|
|
|
# This is the public interface,
|
|
def insert_into_model(model_class: Type, all_data: dict | list, override=False) -> DatabaseOperation:
|
|
"""
|
|
Purpose:
|
|
Generic ORM insert for any model. Keeping it generic for reuse.
|
|
|
|
Args:
|
|
model_class: The ORM model class
|
|
all_data: Dictionary or list of dictionaries with data matching the model
|
|
override: If True, wipe the table before inserting
|
|
|
|
Returns:
|
|
DatabaseOperation object with true/false
|
|
"""
|
|
# Normalize to list for uniform handling
|
|
data_list = all_data if isinstance(all_data, list) else [all_data]
|
|
|
|
# Call the wrapped function with normalized data
|
|
return _wrapped_insert(model_class=model_class, data_list=data_list, override=override)
|
|
|
|
|
|
@safe_db_operation
|
|
def _wrapped_insert(model_class: Type, data_list: dict | list, session: Session, override=False) -> DatabaseOperation:
|
|
"""Insert items. Doesn't manage session. It's managed by the wrapper safe_db_operation"""
|
|
|
|
if override:
|
|
logger.info(f"First, WIPING the pre-existing data for {model_class.__name__}. Are you sure you intended to completely delete the old data?")
|
|
session.query(model_class).delete()
|
|
|
|
logger.info(f"Starting insert for {model_class.__name__}")
|
|
|
|
# Get valid field names upfront
|
|
valid_fields = {col.name for col in model_class.__table__.columns}
|
|
|
|
# Filter data proactively
|
|
filtered_data_list = []
|
|
all_dropped_fields = set()
|
|
|
|
for each_json in data_list:
|
|
filtered_json = {k: v for k, v in each_json.items() if k in valid_fields}
|
|
filtered_data_list.append(filtered_json)
|
|
|
|
# Track dropped fields
|
|
dropped = set(each_json.keys()) - valid_fields
|
|
all_dropped_fields.update(dropped)
|
|
|
|
# Log if we dropped anything
|
|
if all_dropped_fields:
|
|
logger.warning(f"Filtered unknown fields before insert: {all_dropped_fields}")
|
|
|
|
try:
|
|
for filtered_json in filtered_data_list:
|
|
instance = model_class(**filtered_json)
|
|
session.add(instance)
|
|
session.commit()
|
|
|
|
logger.info("Finished wrapped insert!")
|
|
return DatabaseOperation(
|
|
valid=True,
|
|
tried_filtered=bool(all_dropped_fields),
|
|
message=f"Inserted successfully" + (f" after filtering {all_dropped_fields}" if all_dropped_fields else "")
|
|
)
|
|
|
|
except Exception as e:
|
|
# Unexpected error during insert
|
|
session.rollback()
|
|
raise WrapperRollback(DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.UNEXPECTED_DB_ERROR,
|
|
message=f"Wrapped Insert Failed: {str(e)}"
|
|
))
|
|
|
|
# =================== Isolation Test ======================
|
|
"""
|
|
Spoof Failure of an operational error to trigger/test a migration
|
|
"""
|
|
# from sqlalchemy import text
|
|
|
|
# # 1. Drop a column the model expects
|
|
# if model_class == Location:
|
|
# with get_session() as session:
|
|
# print("dropping provider_name from location..")
|
|
# session.execute(text("ALTER TABLE locations DROP COLUMN provider_name"))
|
|
# session.commit()
|
|
# print("dropped provider_name from location!!!")
|
|
|
|
# # 2. Now insert will trigger OperationalError
|