Skip to content

Commit 5c16cc3

Browse files
jirivranyjirivrany
committed
add test for base migration
1 parent 33d68ed commit 5c16cc3

2 files changed

Lines changed: 1450 additions & 90 deletions

File tree

‎flowapp/migrations/versions/001_baseline.py‎

Lines changed: 21 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
Create Date: 2026-02-13
1515
1616
"""
17+
1718
from alembic import op
1819
import sqlalchemy as sa
1920

@@ -76,16 +77,12 @@ def upgrade():
7677
sa.Column("limit_flowspec6", sa.Integer(), default=0),
7778
sa.Column("limit_rtbh", sa.Integer(), default=0),
7879
)
79-
_seed_orgs = True
8080
else:
8181
organization_table = None
82-
_seed_orgs = False
8382
# Add limit columns if missing (pre-v1.0 databases)
8483
for col_name in ("limit_flowspec4", "limit_flowspec6", "limit_rtbh"):
8584
if not _column_exists("organization", col_name):
86-
op.add_column(
87-
"organization", sa.Column(col_name, sa.Integer(), default=0)
88-
)
85+
op.add_column("organization", sa.Column(col_name, sa.Integer(), default=0))
8986

9087
if not _table_exists("rstate"):
9188
rstate_table = op.create_table(
@@ -144,21 +141,15 @@ def upgrade():
144141
if not _table_exists("user_role"):
145142
op.create_table(
146143
"user_role",
147-
sa.Column(
148-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
149-
),
150-
sa.Column(
151-
"role_id", sa.Integer(), sa.ForeignKey("role.id"), nullable=False
152-
),
144+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
145+
sa.Column("role_id", sa.Integer(), sa.ForeignKey("role.id"), nullable=False),
153146
sa.PrimaryKeyConstraint("user_id", "role_id"),
154147
)
155148

156149
if not _table_exists("user_organization"):
157150
op.create_table(
158151
"user_organization",
159-
sa.Column(
160-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
161-
),
152+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
162153
sa.Column(
163154
"organization_id",
164155
sa.Integer(),
@@ -177,9 +168,7 @@ def upgrade():
177168
sa.Column("name", sa.String(length=120), unique=True),
178169
sa.Column("command", sa.String(length=120), unique=True),
179170
sa.Column("description", sa.String(length=260)),
180-
sa.Column(
181-
"role_id", sa.Integer(), sa.ForeignKey("role.id"), nullable=False
182-
),
171+
sa.Column("role_id", sa.Integer(), sa.ForeignKey("role.id"), nullable=False),
183172
)
184173
_seed_actions = True
185174
else:
@@ -202,9 +191,7 @@ def upgrade():
202191
sa.Column("extcomm", sa.String(length=2047)),
203192
sa.Column("description", sa.String(length=255)),
204193
sa.Column("as_path", sa.Boolean(), default=False),
205-
sa.Column(
206-
"role_id", sa.Integer(), sa.ForeignKey("role.id"), nullable=False
207-
),
194+
sa.Column("role_id", sa.Integer(), sa.ForeignKey("role.id"), nullable=False),
208195
)
209196
_seed_communities = True
210197
else:
@@ -244,9 +231,7 @@ def upgrade():
244231
sa.Column("readonly", sa.Boolean(), default=False),
245232
sa.Column("expires", sa.DateTime(), nullable=True),
246233
sa.Column("comment", sa.String(length=255)),
247-
sa.Column(
248-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
249-
),
234+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
250235
sa.Column(
251236
"org_id",
252237
sa.Integer(),
@@ -272,7 +257,6 @@ def upgrade():
272257
sa.Column(
273258
"org_id",
274259
sa.Integer(),
275-
sa.ForeignKey("organization.id"),
276260
nullable=True,
277261
server_default="0",
278262
),
@@ -287,9 +271,7 @@ def upgrade():
287271
sa.Column("readonly", sa.Boolean(), default=True),
288272
sa.Column("expires", sa.DateTime(), nullable=True),
289273
sa.Column("comment", sa.String(length=255)),
290-
sa.Column(
291-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
292-
),
274+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
293275
sa.Column(
294276
"org_id",
295277
sa.Integer(),
@@ -305,7 +287,6 @@ def upgrade():
305287
sa.Column(
306288
"org_id",
307289
sa.Integer(),
308-
sa.ForeignKey("organization.id"),
309290
nullable=True,
310291
),
311292
)
@@ -329,21 +310,15 @@ def upgrade():
329310
sa.Column("comment", sa.Text()),
330311
sa.Column("expires", sa.DateTime()),
331312
sa.Column("created", sa.DateTime()),
332-
sa.Column(
333-
"action_id", sa.Integer(), sa.ForeignKey("action.id"), nullable=False
334-
),
335-
sa.Column(
336-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
337-
),
313+
sa.Column("action_id", sa.Integer(), sa.ForeignKey("action.id"), nullable=False),
314+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
338315
sa.Column(
339316
"org_id",
340317
sa.Integer(),
341318
sa.ForeignKey("organization.id"),
342319
nullable=False,
343320
),
344-
sa.Column(
345-
"rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False
346-
),
321+
sa.Column("rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False),
347322
)
348323
else:
349324
if not _column_exists("flowspec4", "fragment"):
@@ -357,7 +332,6 @@ def upgrade():
357332
sa.Column(
358333
"org_id",
359334
sa.Integer(),
360-
sa.ForeignKey("organization.id"),
361335
nullable=True,
362336
),
363337
)
@@ -378,21 +352,15 @@ def upgrade():
378352
sa.Column("comment", sa.Text()),
379353
sa.Column("expires", sa.DateTime()),
380354
sa.Column("created", sa.DateTime()),
381-
sa.Column(
382-
"action_id", sa.Integer(), sa.ForeignKey("action.id"), nullable=False
383-
),
384-
sa.Column(
385-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
386-
),
355+
sa.Column("action_id", sa.Integer(), sa.ForeignKey("action.id"), nullable=False),
356+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
387357
sa.Column(
388358
"org_id",
389359
sa.Integer(),
390360
sa.ForeignKey("organization.id"),
391361
nullable=False,
392362
),
393-
sa.Column(
394-
"rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False
395-
),
363+
sa.Column("rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False),
396364
)
397365
else:
398366
if not _column_exists("flowspec6", "org_id"):
@@ -401,7 +369,6 @@ def upgrade():
401369
sa.Column(
402370
"org_id",
403371
sa.Integer(),
404-
sa.ForeignKey("organization.id"),
405372
nullable=True,
406373
),
407374
)
@@ -423,18 +390,14 @@ def upgrade():
423390
sa.Column("comment", sa.Text()),
424391
sa.Column("expires", sa.DateTime()),
425392
sa.Column("created", sa.DateTime()),
426-
sa.Column(
427-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
428-
),
393+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
429394
sa.Column(
430395
"org_id",
431396
sa.Integer(),
432397
sa.ForeignKey("organization.id"),
433398
nullable=False,
434399
),
435-
sa.Column(
436-
"rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False
437-
),
400+
sa.Column("rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False),
438401
)
439402
else:
440403
if not _column_exists("RTBH", "org_id"):
@@ -443,7 +406,6 @@ def upgrade():
443406
sa.Column(
444407
"org_id",
445408
sa.Integer(),
446-
sa.ForeignKey("organization.id"),
447409
nullable=True,
448410
),
449411
)
@@ -457,18 +419,14 @@ def upgrade():
457419
sa.Column("comment", sa.Text()),
458420
sa.Column("expires", sa.DateTime()),
459421
sa.Column("created", sa.DateTime()),
460-
sa.Column(
461-
"user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False
462-
),
422+
sa.Column("user_id", sa.Integer(), sa.ForeignKey("user.id"), nullable=False),
463423
sa.Column(
464424
"org_id",
465425
sa.Integer(),
466426
sa.ForeignKey("organization.id"),
467427
nullable=False,
468428
),
469-
sa.Column(
470-
"rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False
471-
),
429+
sa.Column("rstate_id", sa.Integer(), sa.ForeignKey("rstate.id"), nullable=False),
472430
)
473431

474432
if not _table_exists("rule_whitelist_cache"):
@@ -498,39 +456,12 @@ def upgrade():
498456
],
499457
)
500458

501-
if _seed_orgs and organization_table is not None:
502-
op.bulk_insert(
503-
organization_table,
504-
[
505-
{
506-
"name": "TU Liberec",
507-
"arange": "147.230.0.0/16\n2001:718:1c01::/48",
508-
"limit_flowspec4": 0,
509-
"limit_flowspec6": 0,
510-
"limit_rtbh": 0,
511-
},
512-
{
513-
"name": "Cesnet",
514-
"arange": "147.230.0.0/16\n2001:718:1c01::/48",
515-
"limit_flowspec4": 0,
516-
"limit_flowspec6": 0,
517-
"limit_rtbh": 0,
518-
},
519-
],
520-
)
521-
522459
# Ensure rstate has the "whitelisted rule" entry (id=4, added in v1.1.0)
523460
if not _seed_rstates and _table_has_data("rstate"):
524461
conn = op.get_bind()
525-
result = conn.execute(
526-
sa.text("SELECT COUNT(*) FROM rstate WHERE id = 4")
527-
)
462+
result = conn.execute(sa.text("SELECT COUNT(*) FROM rstate WHERE id = 4"))
528463
if result.scalar() == 0:
529-
conn.execute(
530-
sa.text(
531-
"INSERT INTO rstate (id, description) VALUES (4, 'whitelisted rule')"
532-
)
533-
)
464+
conn.execute(sa.text("INSERT INTO rstate (id, description) VALUES (4, 'whitelisted rule')"))
534465

535466
if _seed_rstates and not _table_has_data("rstate"):
536467
op.bulk_insert(

0 commit comments

Comments
 (0)