115 lines
4.3 KiB
Python
115 lines
4.3 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
|
|
|
|
# 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.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
|
|
"""
|
|
logger.info(f"All the public data preparing to be inserted is {all_data}")
|
|
|
|
# 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."""
|
|
|
|
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__}")
|
|
|
|
try:
|
|
for each_json in data_list:
|
|
instance = model_class(**each_json)
|
|
session.add(instance)
|
|
session.commit() # Commits both delete + inserts atomically
|
|
return DatabaseOperation(valid=True)
|
|
|
|
except TypeError as e:
|
|
"""
|
|
This is triggered when an Old client was sent unknown fields from the API in the JSON,
|
|
And it's handled inside the wrapped function because it needs access to the variables used.
|
|
"""
|
|
session.rollback()
|
|
|
|
valid_fields = {col.name for col in model_class.__table__.columns}
|
|
filtered_data_list = [
|
|
{k: v for k, v in each_json.items() if k in valid_fields}
|
|
for each_json in data_list
|
|
]
|
|
|
|
dropped_fields = set().union(*(set(d.keys()) - valid_fields for d in data_list))
|
|
if dropped_fields:
|
|
logger.warning(f"Old Client Dropped unknown fields: {dropped_fields}")
|
|
|
|
try:
|
|
for filtered_json in filtered_data_list:
|
|
instance = model_class(**filtered_json)
|
|
session.add(instance)
|
|
session.commit()
|
|
|
|
# the fix worked,
|
|
return DatabaseOperation(
|
|
valid=True,
|
|
tried_filtered=True,
|
|
message=f"Inserted after filtering {dropped_fields}"
|
|
)
|
|
|
|
except Exception as retry_error:
|
|
# Signal wrapper to rollback
|
|
raise WrapperRollback(DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.OLD_CLIENT_NEW_API,
|
|
tried_filtered=True,
|
|
message=str(f"Wrapped Insert Failed on Retry of a TypeError: {retry_error}")
|
|
))
|
|
|
|
|
|
# =================== 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
|