Skip to content

Commit

Permalink
Do not support MSSQL less than v2017 in code
Browse files Browse the repository at this point in the history
Our experimental support for MSSQL starts from v2017(in README.md) but
we still support 2000 & 2005 in code.
This PR removes this support, allowing us to use mssql.DATETIME2 in all
MSSQL DB.
  • Loading branch information
ephraimbuddy committed Jun 2, 2022
1 parent 1c680e4 commit dcdfce2
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 40 deletions.
25 changes: 4 additions & 21 deletions airflow/migrations/db_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,31 +29,14 @@
######################################


def _mssql_use_date_time2():
conn = context.get_bind()
result = conn.execute(
"""SELECT CASE WHEN CONVERT(VARCHAR(128), SERVERPROPERTY ('productversion'))
like '8%' THEN '2000' WHEN CONVERT(VARCHAR(128), SERVERPROPERTY ('productversion'))
like '9%' THEN '2005' ELSE '2005Plus' END AS MajorVersion"""
).fetchone()
mssql_version = result[0]
return mssql_version not in ("2000", "2005")


MSSQL_USE_DATE_TIME2 = Proxy(_mssql_use_date_time2)


def _mssql_TIMESTAMP():
from sqlalchemy.dialects import mssql

if MSSQL_USE_DATE_TIME2:

class DATETIME2(mssql.DATETIME2):
def __init__(self, *args, precision=6, **kwargs):
super().__init__(*args, precision=precision, **kwargs)
class DATETIME2(mssql.DATETIME2):
def __init__(self, *args, precision=6, **kwargs):
super().__init__(*args, precision=precision, **kwargs)

return DATETIME2
return mssql.DATETIME
return DATETIME2


def _mysql_TIMESTAMP():
Expand Down
2 changes: 0 additions & 2 deletions airflow/migrations/db_types.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -24,5 +24,3 @@ TIMESTAMP = sa.TIMESTAMP

StringID = sa.String
"""String column type with correct DB collation applied"""

MSSQL_USE_DATE_TIME2: bool
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
from alembic import op
from sqlalchemy.dialects import mssql

from airflow.migrations.db_types import MSSQL_USE_DATE_TIME2, TIMESTAMP
from airflow.migrations.db_types import TIMESTAMP

# revision identifiers, used by Alembic.
revision = '83f031fd9f1c'
Expand Down Expand Up @@ -135,21 +135,12 @@ def recreate_mssql_ts_column(conn, op, table_name, column_name):

def alter_mssql_datetime_column(conn, op, table_name, column_name, nullable):
"""Update the datetime column to datetime2(6)"""
if MSSQL_USE_DATE_TIME2:
op.alter_column(
table_name=table_name,
column_name=column_name,
type_=mssql.DATETIME2(precision=6),
nullable=nullable,
)


def alter_mssql_datetime2_column(conn, op, table_name, column_name, nullable):
"""Update the datetime2(6) column to datetime"""
if MSSQL_USE_DATE_TIME2:
op.alter_column(
table_name=table_name, column_name=column_name, type_=mssql.DATETIME, nullable=nullable
)
op.alter_column(
table_name=table_name,
column_name=column_name,
type_=mssql.DATETIME2(precision=6),
nullable=nullable,
)


def upgrade():
Expand Down Expand Up @@ -199,7 +190,6 @@ def downgrade():
conn = op.get_bind()
if conn.dialect.name != 'mssql':
return
alter_mssql_datetime2_column(conn, op, 'serialized_dag', 'last_updated', False)
op.alter_column(table_name="xcom", column_name="timestamp", type_=TIMESTAMP, nullable=True)
with op.batch_alter_table('task_reschedule') as task_reschedule_batch_op:
task_reschedule_batch_op.alter_column(column_name='end_date', type_=TIMESTAMP, nullable=True)
Expand Down

0 comments on commit dcdfce2

Please sign in to comment.