Skip to content

Commit

Permalink
Do not support MSSQL less than v2017 in code (#24095)
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 authored Jun 2, 2022
1 parent f2e6452 commit 75fdbf0
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 75fdbf0

Please sign in to comment.