Skip to content

Commit

Permalink
Add check constraints for sqlite
Browse files Browse the repository at this point in the history
  • Loading branch information
dhirving committed Mar 14, 2024
1 parent e7bedaf commit 0db0a35
Showing 1 changed file with 17 additions and 0 deletions.
17 changes: 17 additions & 0 deletions migrations/dimensions-config/1fae088c80b6.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,9 @@ def _lock_exposure_table(ctx: _Context) -> None:

def _migrate_groups(ctx: _Context) -> None:
# Create group table
check_constraints = []
if ctx.is_sqlite:
check_constraints = [_make_string_check_constraint("instrument", 32, "group_len_instrument")]
table = op.create_table(
"group",
sa.Column("instrument", sa.String(32), primary_key=True),
Expand All @@ -70,6 +73,7 @@ def _migrate_groups(ctx: _Context) -> None:
refcolumns=["instrument.name"],
name="fkey_group_instrument_name_instrument",
),
*check_constraints,
schema=ctx.schema,
)

Expand Down Expand Up @@ -127,6 +131,10 @@ def _migrate_day_obs(ctx: _Context) -> None:
sa.Column("timespan_end", sa.BigInteger),
]

check_constraints = []
if ctx.is_sqlite:
check_constraints = [_make_string_check_constraint("instrument", 32, "day_obs_len_instrument")]

table = op.create_table(
"day_obs",
sa.Column("instrument", sa.String(32), primary_key=True),
Expand All @@ -137,6 +145,7 @@ def _migrate_day_obs(ctx: _Context) -> None:
refcolumns=["instrument.name"],
name="fkey_day_obs_instrument_name_instrument",
),
*check_constraints,
schema=ctx.schema,
)

Expand Down Expand Up @@ -204,6 +213,14 @@ def _get_day_obs_offset(instrument_name: str, instrument: _Instrument, day_obs:
return round(offset.to_value("s"))


def _make_string_check_constraint(
field_name: str, max_length: int, constraint_name: str
) -> sa.schema.CheckConstraint:
return sa.schema.CheckConstraint(
f'length("{field_name}")<={max_length} AND length("{field_name}")>=1', name=constraint_name
)


class _Context:
def __init__(self) -> None:
self.mig_context = alembic.context.get_context()
Expand Down

0 comments on commit 0db0a35

Please sign in to comment.