Skip to content

Commit

Permalink
Fix for sqlalchemy 2 and text upgrades
Browse files Browse the repository at this point in the history
  • Loading branch information
pjaol committed Dec 20, 2023
1 parent b5be147 commit 6e85490
Showing 1 changed file with 11 additions and 8 deletions.
19 changes: 11 additions & 8 deletions luigi/db_task_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@
import sqlalchemy.orm
import sqlalchemy.orm.collections
from sqlalchemy.engine import reflection
from sqlalchemy import text

Base = sqlalchemy.ext.declarative.declarative_base()

logger = logging.getLogger('luigi-interface')
Expand Down Expand Up @@ -251,25 +253,26 @@ def _upgrade_schema(engine):
# Upgrade 1. Add task_id column and index to tasks
if 'task_id' not in [x['name'] for x in inspector.get_columns('tasks')]:
logger.warning('Upgrading DbTaskHistory schema: Adding tasks.task_id')
conn.execute('ALTER TABLE tasks ADD COLUMN task_id VARCHAR(200)')
conn.execute('CREATE INDEX ix_task_id ON tasks (task_id)')
conn.execute(text('ALTER TABLE tasks ADD COLUMN task_id VARCHAR(200)'))
conn.execute(text('CREATE INDEX ix_task_id ON tasks (task_id)'))

Check warning on line 257 in luigi/db_task_history.py

View check run for this annotation

Codecov / codecov/patch

luigi/db_task_history.py#L256-L257

Added lines #L256 - L257 were not covered by tests

# Upgrade 2. Alter value column to be TEXT, note that this is idempotent so no if-guard
if 'mysql' in engine.dialect.name:
conn.execute('ALTER TABLE task_parameters MODIFY COLUMN value TEXT')
conn.execute(text('ALTER TABLE task_parameters MODIFY COLUMN value TEXT'))

Check warning on line 261 in luigi/db_task_history.py

View check run for this annotation

Codecov / codecov/patch

luigi/db_task_history.py#L261

Added line #L261 was not covered by tests
elif 'oracle' in engine.dialect.name:
conn.execute('ALTER TABLE task_parameters MODIFY value TEXT')
conn.execute(text('ALTER TABLE task_parameters MODIFY value TEXT'))

Check warning on line 263 in luigi/db_task_history.py

View check run for this annotation

Codecov / codecov/patch

luigi/db_task_history.py#L263

Added line #L263 was not covered by tests
elif 'mssql' in engine.dialect.name:
conn.execute('ALTER TABLE task_parameters ALTER COLUMN value TEXT')
conn.execute(text('ALTER TABLE task_parameters ALTER COLUMN value TEXT'))

Check warning on line 265 in luigi/db_task_history.py

View check run for this annotation

Codecov / codecov/patch

luigi/db_task_history.py#L265

Added line #L265 was not covered by tests
elif 'postgresql' in engine.dialect.name:
if str([x for x in inspector.get_columns('task_parameters')
if x['name'] == 'value'][0]['type']) != 'TEXT':
conn.execute('ALTER TABLE task_parameters ALTER COLUMN value TYPE TEXT')
conn.execute(text('ALTER TABLE task_parameters ALTER COLUMN value TYPE TEXT'))

Check warning on line 269 in luigi/db_task_history.py

View check run for this annotation

Codecov / codecov/patch

luigi/db_task_history.py#L269

Added line #L269 was not covered by tests
elif 'sqlite' in engine.dialect.name:
# SQLite does not support changing column types. A database file will need
# to be used to pickup this migration change.
for i in conn.execute('PRAGMA table_info(task_parameters);').fetchall():
if i['name'] == 'value' and i['type'] != 'TEXT':
for i in conn.execute(text('PRAGMA table_info(task_parameters);')).fetchall():
x = i._asdict()
if x['name'] == 'value' and x['type'] != 'TEXT':

Check warning on line 275 in luigi/db_task_history.py

View check run for this annotation

Codecov / codecov/patch

luigi/db_task_history.py#L273-L275

Added lines #L273 - L275 were not covered by tests
logger.warning(
'SQLite can not change column types. Please use a new database '
'to pickup column type changes.'
Expand Down

0 comments on commit 6e85490

Please sign in to comment.