Skip to content

Commit ddbd430

Browse files
LennartSchmidtKernJohannes HötterJWittmeyer
authored
Smart Planner Unprototype (#248)
* adds iteration number to pipeline logs * adds strategy requirements and mapping options * is input to database model * merge dev * alembic rehead * routing config * model * model * union alembic * rename alembic * alembic update * model * Alembic fix --------- Co-authored-by: Johannes Hötter <[email protected]> Co-authored-by: JWittmeyer <[email protected]>
1 parent ec7349b commit ddbd430

File tree

2 files changed

+194
-1
lines changed

2 files changed

+194
-1
lines changed
Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
"""Adds smart routing
2+
3+
Revision ID: 3e59ce51739c
4+
Revises: b7b9f7de1013
5+
Create Date: 2024-09-02 12:32:23.493855
6+
7+
"""
8+
9+
from alembic import op
10+
import sqlalchemy as sa
11+
from sqlalchemy.dialects import postgresql
12+
13+
# revision identifiers, used by Alembic.
14+
revision = "3e59ce51739c"
15+
down_revision = "b7b9f7de1013"
16+
branch_labels = None
17+
depends_on = None
18+
19+
20+
def upgrade():
21+
22+
connection = op.get_bind()
23+
24+
# ### commands auto generated by Alembic - please adjust! ###
25+
op.create_table(
26+
"strategy_requirement",
27+
sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False),
28+
sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=True),
29+
sa.Column("strategy_id", postgresql.UUID(as_uuid=True), nullable=True),
30+
sa.Column("created_by", postgresql.UUID(as_uuid=True), nullable=True),
31+
sa.Column("created_at", sa.DateTime(), nullable=True),
32+
sa.Column("field", sa.String(), nullable=True),
33+
sa.Column("description", sa.String(), nullable=True),
34+
sa.Column("is_input", sa.Boolean(), nullable=True),
35+
sa.ForeignKeyConstraint(["created_by"], ["user.id"], ondelete="SET NULL"),
36+
sa.ForeignKeyConstraint(
37+
["project_id"], ["cognition.project.id"], ondelete="CASCADE"
38+
),
39+
sa.ForeignKeyConstraint(
40+
["strategy_id"], ["cognition.strategy.id"], ondelete="CASCADE"
41+
),
42+
sa.PrimaryKeyConstraint("id"),
43+
schema="cognition",
44+
)
45+
op.create_index(
46+
op.f("ix_cognition_strategy_requirement_created_by"),
47+
"strategy_requirement",
48+
["created_by"],
49+
unique=False,
50+
schema="cognition",
51+
)
52+
op.create_index(
53+
op.f("ix_cognition_strategy_requirement_project_id"),
54+
"strategy_requirement",
55+
["project_id"],
56+
unique=False,
57+
schema="cognition",
58+
)
59+
op.create_index(
60+
op.f("ix_cognition_strategy_requirement_strategy_id"),
61+
"strategy_requirement",
62+
["strategy_id"],
63+
unique=False,
64+
schema="cognition",
65+
)
66+
op.create_table(
67+
"strategy_requirement_mapping_option",
68+
sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False),
69+
sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=True),
70+
sa.Column(
71+
"strategy_requirement_id", postgresql.UUID(as_uuid=True), nullable=True
72+
),
73+
sa.Column("created_by", postgresql.UUID(as_uuid=True), nullable=True),
74+
sa.Column("created_at", sa.DateTime(), nullable=True),
75+
sa.Column("value", sa.String(), nullable=True),
76+
sa.ForeignKeyConstraint(["created_by"], ["user.id"], ondelete="SET NULL"),
77+
sa.ForeignKeyConstraint(
78+
["project_id"], ["cognition.project.id"], ondelete="CASCADE"
79+
),
80+
sa.ForeignKeyConstraint(
81+
["strategy_requirement_id"],
82+
["cognition.strategy_requirement.id"],
83+
ondelete="CASCADE",
84+
),
85+
sa.PrimaryKeyConstraint("id"),
86+
schema="cognition",
87+
)
88+
op.create_index(
89+
op.f("ix_cognition_strategy_requirement_mapping_option_created_by"),
90+
"strategy_requirement_mapping_option",
91+
["created_by"],
92+
unique=False,
93+
schema="cognition",
94+
)
95+
op.create_index(
96+
op.f("ix_cognition_strategy_requirement_mapping_option_project_id"),
97+
"strategy_requirement_mapping_option",
98+
["project_id"],
99+
unique=False,
100+
schema="cognition",
101+
)
102+
op.create_index(
103+
op.f(
104+
"ix_cognition_strategy_requirement_mapping_option_strategy_requirement_id"
105+
),
106+
"strategy_requirement_mapping_option",
107+
["strategy_requirement_id"],
108+
unique=False,
109+
schema="cognition",
110+
)
111+
op.add_column(
112+
"pipeline_logs",
113+
sa.Column("iteration_number", sa.Integer(), nullable=True),
114+
schema="cognition",
115+
)
116+
op.add_column(
117+
"project",
118+
sa.Column("operator_routing_config", sa.JSON(), nullable=True),
119+
schema="cognition",
120+
)
121+
122+
update_project_sql = """
123+
UPDATE cognition.project
124+
SET operator_routing_config = jsonb_build_object(
125+
'sourceCode', operator_routing_source_code::text
126+
)
127+
WHERE operator_routing_source_code IS NOT NULL
128+
"""
129+
connection.execute(update_project_sql)
130+
131+
op.drop_column("project", "operator_routing_source_code", schema="cognition")
132+
# ### end Alembic commands ###
133+
134+
135+
def downgrade():
136+
connection = op.get_bind()
137+
# ### commands auto generated by Alembic - please adjust! ###
138+
op.add_column(
139+
"project",
140+
sa.Column(
141+
"operator_routing_source_code",
142+
sa.VARCHAR(),
143+
autoincrement=False,
144+
nullable=True,
145+
),
146+
schema="cognition",
147+
)
148+
149+
reverse_update_project_sql = """
150+
UPDATE cognition.project
151+
SET operator_routing_source_code = (operator_routing_config->>'sourceCode')
152+
WHERE operator_routing_config IS NOT NULL
153+
AND operator_routing_config->>'sourceCode' IS NOT NULL;
154+
"""
155+
connection.execute(reverse_update_project_sql)
156+
157+
op.drop_column("project", "operator_routing_config", schema="cognition")
158+
op.drop_column("pipeline_logs", "iteration_number", schema="cognition")
159+
op.drop_index(
160+
op.f(
161+
"ix_cognition_strategy_requirement_mapping_option_strategy_requirement_id"
162+
),
163+
table_name="strategy_requirement_mapping_option",
164+
schema="cognition",
165+
)
166+
op.drop_index(
167+
op.f("ix_cognition_strategy_requirement_mapping_option_project_id"),
168+
table_name="strategy_requirement_mapping_option",
169+
schema="cognition",
170+
)
171+
op.drop_index(
172+
op.f("ix_cognition_strategy_requirement_mapping_option_created_by"),
173+
table_name="strategy_requirement_mapping_option",
174+
schema="cognition",
175+
)
176+
op.drop_table("strategy_requirement_mapping_option", schema="cognition")
177+
op.drop_index(
178+
op.f("ix_cognition_strategy_requirement_strategy_id"),
179+
table_name="strategy_requirement",
180+
schema="cognition",
181+
)
182+
op.drop_index(
183+
op.f("ix_cognition_strategy_requirement_project_id"),
184+
table_name="strategy_requirement",
185+
schema="cognition",
186+
)
187+
op.drop_index(
188+
op.f("ix_cognition_strategy_requirement_created_by"),
189+
table_name="strategy_requirement",
190+
schema="cognition",
191+
)
192+
op.drop_table("strategy_requirement", schema="cognition")
193+
# ### end Alembic commands ###

0 commit comments

Comments
 (0)