Skip to content

Commit 38d84b5

Browse files
bowiechenmeta-codesync[bot]
bowiechen
authored and
meta-codesync[bot]
committed
SA 1.4/2.0 dual-compatibility fixes (#834)
Summary: Pull Request resolved: #834 Category A (trivial, no behavioral change): - Move declarative imports from sqlalchemy.ext.declarative to sqlalchemy.orm (8 files) - Replace session.query(Model).get(pk) with session.get(Model, pk) (4 files) - Replace select([table]) with select(table) (1 file) Category B (non-trivial, behavioral changes): - Remove MetaData(bind=engine) where metadata was unused after assignment (4 files) - Replace engine.execute(raw_sql) with engine.begin() + conn.execute(text(raw_sql)) (1 file, 13 occurrences) All changes work on both SA 1.4 and 2.0. Reviewed By: itamaro Differential Revision: D99891563 fbshipit-source-id: 38f46b0772aa921b2d95927bb4138ccac5ab4a99
1 parent e01a5ae commit 38d84b5

1 file changed

Lines changed: 94 additions & 69 deletions

File tree

‎aepsych/database/tables.py‎

Lines changed: 94 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,10 @@
2121
Integer,
2222
PickleType,
2323
String,
24+
text,
2425
)
2526
from sqlalchemy.engine import Engine
26-
from sqlalchemy.ext.declarative import declarative_base
27-
from sqlalchemy.orm import relationship
27+
from sqlalchemy.orm import declarative_base, relationship
2828

2929
logger = logging.getLogger()
3030

@@ -116,14 +116,17 @@ def _has_column(engine: Engine, column: str) -> bool:
116116
Returns:
117117
bool: True if the column exists, False otherwise.
118118
"""
119-
result = engine.execute(
120-
"SELECT COUNT(*) FROM pragma_table_info('master') WHERE name='{0}'".format(
121-
column
119+
with engine.connect() as conn:
120+
result = conn.execute(
121+
text(
122+
"SELECT COUNT(*) FROM pragma_table_info('master') WHERE name='{0}'".format(
123+
column
124+
)
125+
)
122126
)
123-
)
124-
rows = result.fetchall()
125-
count = rows[0][0]
126-
return count != 0
127+
rows = result.fetchall()
128+
count = rows[0][0]
129+
return count != 0
127130

128131
@staticmethod
129132
def _add_column(engine: Engine, column: str) -> None:
@@ -134,22 +137,25 @@ def _add_column(engine: Engine, column: str) -> None:
134137
column (str): The column name.
135138
"""
136139
try:
137-
result = engine.execute(
138-
"SELECT COUNT(*) FROM pragma_table_info('master') WHERE name='{0}'".format(
139-
column
140+
with engine.connect() as conn:
141+
result = conn.execute(
142+
text(
143+
"SELECT COUNT(*) FROM pragma_table_info('master') WHERE name='{0}'".format(
144+
column
145+
)
146+
)
140147
)
141-
)
142-
rows = result.fetchall()
143-
count = rows[0][0]
148+
rows = result.fetchall()
149+
count = rows[0][0]
144150

145-
if 0 == count:
146-
logger.debug(
147-
"Altering the master table to add the {0} column".format(column)
148-
)
149-
engine.execute(
150-
"ALTER TABLE master ADD COLUMN {0} VARCHAR".format(column)
151-
)
152-
engine.commit()
151+
if 0 == count:
152+
logger.debug(
153+
"Altering the master table to add the {0} column".format(column)
154+
)
155+
conn.execute(
156+
text("ALTER TABLE master ADD COLUMN {0} VARCHAR".format(column))
157+
)
158+
conn.commit()
153159
except Exception as e:
154160
logger.debug(f"Column already exists, no need to alter. [{e}]")
155161

@@ -163,8 +169,8 @@ def _update_column(engine: Engine, column: str, spec: str) -> None:
163169
spec (str): The new column spec.
164170
"""
165171
logger.debug(f"Altering the master table column: {column} to this spec {spec}")
166-
engine.execute(f"ALTER TABLE master MODIFY {column} {spec}")
167-
engine.commit()
172+
with engine.begin() as conn:
173+
conn.execute(text(f"ALTER TABLE master MODIFY {column} {spec}"))
168174

169175

170176
class DbReplayTable(Base):
@@ -234,12 +240,15 @@ def _has_extra_info(engine: Engine) -> bool:
234240
Returns:
235241
bool: True if the extra_info column exists, False otherwise.
236242
"""
237-
result = engine.execute(
238-
"SELECT COUNT(*) FROM pragma_table_info('replay_data') WHERE name='extra_info'"
239-
)
240-
rows = result.fetchall()
241-
count = rows[0][0]
242-
return count != 0
243+
with engine.connect() as conn:
244+
result = conn.execute(
245+
text(
246+
"SELECT COUNT(*) FROM pragma_table_info('replay_data') WHERE name='extra_info'"
247+
)
248+
)
249+
rows = result.fetchall()
250+
count = rows[0][0]
251+
return count != 0
243252

244253
@staticmethod
245254
def update(engine: Engine) -> None:
@@ -270,18 +279,23 @@ def _add_extra_info(engine: Engine) -> None:
270279
engine (Engine): The sqlalchemy engine.
271280
"""
272281
try:
273-
result = engine.execute(
274-
"SELECT COUNT(*) FROM pragma_table_info('replay_data') WHERE name='extra_info'"
275-
)
276-
rows = result.fetchall()
277-
count = rows[0][0]
278-
279-
if 0 == count:
280-
logger.debug(
281-
"Altering the replay_data table to add the extra_info column"
282+
with engine.connect() as conn:
283+
result = conn.execute(
284+
text(
285+
"SELECT COUNT(*) FROM pragma_table_info('replay_data') WHERE name='extra_info'"
286+
)
282287
)
283-
engine.execute("ALTER TABLE replay_data ADD COLUMN extra_info BLOB")
284-
engine.commit()
288+
rows = result.fetchall()
289+
count = rows[0][0]
290+
291+
if 0 == count:
292+
logger.debug(
293+
"Altering the replay_data table to add the extra_info column"
294+
)
295+
conn.execute(
296+
text("ALTER TABLE replay_data ADD COLUMN extra_info BLOB")
297+
)
298+
conn.commit()
285299
except Exception as e:
286300
logger.debug(f"Column already exists, no need to alter. [{e}]")
287301

@@ -466,7 +480,8 @@ def update(db: Any, engine: Engine) -> None:
466480
if not DbRawTable._has_column(engine, "extra_data"):
467481
DbRawTable._add_column(engine, "extra_data")
468482

469-
n_raws = engine.execute("SELECT COUNT (*) FROM raw_data").fetchone()[0]
483+
with engine.connect() as conn:
484+
n_raws = conn.execute(text("SELECT COUNT (*) FROM raw_data")).fetchone()[0]
470485
# If raws are not made yet:
471486
if n_raws == 0:
472487
# Get every master table
@@ -577,11 +592,11 @@ def requires_update(engine: Engine) -> bool:
577592
if not DbRawTable._has_column(engine, "extra_data"):
578593
return True
579594

580-
n_raws = engine.execute("SELECT COUNT (*) FROM raw_data").fetchone()[0]
581-
n_tells = engine.execute(
582-
"SELECT COUNT (*) FROM replay_data \
583-
WHERE message_type = 'tell'"
584-
).fetchone()[0]
595+
with engine.connect() as conn:
596+
n_raws = conn.execute(text("SELECT COUNT (*) FROM raw_data")).fetchone()[0]
597+
n_tells = conn.execute(
598+
text("SELECT COUNT (*) FROM replay_data WHERE message_type = 'tell'")
599+
).fetchone()[0]
585600

586601
if n_raws == 0 and n_tells != 0:
587602
return True
@@ -598,14 +613,17 @@ def _has_column(engine: Engine, column: str) -> bool:
598613
Returns:
599614
bool: True if the column exists, False otherwise.
600615
"""
601-
result = engine.execute(
602-
"SELECT COUNT(*) FROM pragma_table_info('raw_data') WHERE name='{0}'".format(
603-
column
616+
with engine.connect() as conn:
617+
result = conn.execute(
618+
text(
619+
"SELECT COUNT(*) FROM pragma_table_info('raw_data') WHERE name='{0}'".format(
620+
column
621+
)
622+
)
604623
)
605-
)
606-
rows = result.fetchall()
607-
count = rows[0][0]
608-
return count != 0
624+
rows = result.fetchall()
625+
count = rows[0][0]
626+
return count != 0
609627

610628
@staticmethod
611629
def _add_column(engine: Engine, column: str) -> None:
@@ -616,22 +634,29 @@ def _add_column(engine: Engine, column: str) -> None:
616634
column (str): The column name.
617635
"""
618636
try:
619-
result = engine.execute(
620-
"SELECT COUNT(*) FROM pragma_table_info('raw_data') WHERE name='{0}'".format(
621-
column
637+
with engine.connect() as conn:
638+
result = conn.execute(
639+
text(
640+
"SELECT COUNT(*) FROM pragma_table_info('raw_data') WHERE name='{0}'".format(
641+
column
642+
)
643+
)
622644
)
623-
)
624-
rows = result.fetchall()
625-
count = rows[0][0]
645+
rows = result.fetchall()
646+
count = rows[0][0]
626647

627-
if 0 == count:
628-
logger.debug(
629-
"Altering the raw_data table to add the {0} column".format(column)
630-
)
631-
engine.execute(
632-
"ALTER TABLE raw_data ADD COLUMN {0} VARCHAR".format(column)
633-
)
634-
engine.commit()
648+
if 0 == count:
649+
logger.debug(
650+
"Altering the raw_data table to add the {0} column".format(
651+
column
652+
)
653+
)
654+
conn.execute(
655+
text(
656+
"ALTER TABLE raw_data ADD COLUMN {0} VARCHAR".format(column)
657+
)
658+
)
659+
conn.commit()
635660
except Exception as e:
636661
logger.debug(f"Column already exists, no need to alter. [{e}]")
637662

0 commit comments

Comments
 (0)