def _for_database(cls, database_key: SQLAlchemyDatabaseKey) -> Session:
# TODO(#8046): When the above method is deleted, move this into `using_database`
# directly.
engine = SQLAlchemyEngineManager.get_engine_for_database(
database_key=database_key)
if engine is None:
raise ValueError(f"No engine set for key [{database_key}]")
session = Session(bind=engine)
cls._alter_session_variables(session)
cls._apply_session_listener_for_schema_base(
database_key.declarative_meta, session)
return session
def _write_df_only_successful_rows(table: DeclarativeMeta,
df: pd.DataFrame) -> None:
"""If the dataframe can't be written all at once (eg. some rows already
exist in the database) then we write only the rows that we can."""
for i in range(len(df)):
row = df.iloc[i:i + 1]
try:
row.to_sql(
table.__tablename__,
SQLAlchemyEngineManager.get_engine_for_database(
SQLAlchemyDatabaseKey(schema_type=SchemaType.JAILS)),
if_exists="append",
index=False,
)
except IntegrityError:
# Skip rows that can't be written
logging.info("Skipping write_df to %s table: %s.", table, row)
def write_df(table: DeclarativeMeta, df: pd.DataFrame) -> None:
"""
Writes the |df| to the |table|.
The column headers on |df| must match the column names in |table|. All rows
in |df| will be appended to |table|. If a row in |df| already exists in
|table|, then that row will be skipped.
"""
try:
df.to_sql(
table.__tablename__,
SQLAlchemyEngineManager.get_engine_for_database(
SQLAlchemyDatabaseKey(schema_type=SchemaType.JAILS)),
if_exists="append",
index=False,
)
except IntegrityError:
_write_df_only_successful_rows(table, df)