Skip to content

Commit b8ed2bd

Browse files
jessicasingh7jessicasingh7
committed
accidental removed migrations
1 parent 63bd0f2 commit b8ed2bd

8 files changed

+707
-47
lines changed

‎backend/alembic/versions/5ae8240accb3_add_research_agent_database_tables_and_.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
"""add research agent database tables and chat message research fields
2+
3+
Revision ID: 5ae8240accb3
4+
Revises: b558f51620b4
5+
Create Date: 2025-08-06 14:29:24.691388
6+
7+
"""
8+
9+
from alembic import op
10+
import sqlalchemy as sa
11+
from sqlalchemy.dialects import postgresql
12+
13+
14+
# revision identifiers, used by Alembic.
15+
revision = "5ae8240accb3"
16+
down_revision = "b558f51620b4"
17+
branch_labels = None
18+
depends_on = None
19+
20+
21+
def upgrade() -> None:
22+
# Add research_type and research_plan columns to chat_message table
23+
op.add_column(
24+
"chat_message",
25+
sa.Column("research_type", sa.String(), nullable=True),
26+
)
27+
op.add_column(
28+
"chat_message",
29+
sa.Column("research_plan", postgresql.JSONB(), nullable=True),
30+
)
31+
32+
# Create research_agent_iteration table
33+
op.create_table(
34+
"research_agent_iteration",
35+
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
36+
sa.Column(
37+
"primary_question_id",
38+
sa.Integer(),
39+
sa.ForeignKey("chat_message.id", ondelete="CASCADE"),
40+
nullable=False,
41+
),
42+
sa.Column("iteration_nr", sa.Integer(), nullable=False),
43+
sa.Column(
44+
"created_at",
45+
sa.DateTime(timezone=True),
46+
server_default=sa.func.now(),
47+
nullable=False,
48+
),
49+
sa.Column("purpose", sa.String(), nullable=True),
50+
sa.Column("reasoning", sa.String(), nullable=True),
51+
sa.PrimaryKeyConstraint("id"),
52+
sa.UniqueConstraint(
53+
"primary_question_id",
54+
"iteration_nr",
55+
name="_research_agent_iteration_unique_constraint",
56+
),
57+
)
58+
59+
# Create research_agent_iteration_sub_step table
60+
op.create_table(
61+
"research_agent_iteration_sub_step",
62+
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
63+
sa.Column(
64+
"primary_question_id",
65+
sa.Integer(),
66+
sa.ForeignKey("chat_message.id", ondelete="CASCADE"),
67+
nullable=False,
68+
),
69+
sa.Column(
70+
"parent_question_id",
71+
sa.Integer(),
72+
sa.ForeignKey("research_agent_iteration_sub_step.id", ondelete="CASCADE"),
73+
nullable=True,
74+
),
75+
sa.Column("iteration_nr", sa.Integer(), nullable=False),
76+
sa.Column("iteration_sub_step_nr", sa.Integer(), nullable=False),
77+
sa.Column(
78+
"created_at",
79+
sa.DateTime(timezone=True),
80+
server_default=sa.func.now(),
81+
nullable=False,
82+
),
83+
sa.Column("sub_step_instructions", sa.String(), nullable=True),
84+
sa.Column(
85+
"sub_step_tool_id",
86+
sa.Integer(),
87+
sa.ForeignKey("tool.id"),
88+
nullable=True,
89+
),
90+
sa.Column("reasoning", sa.String(), nullable=True),
91+
sa.Column("sub_answer", sa.String(), nullable=True),
92+
sa.Column("cited_doc_results", postgresql.JSONB(), nullable=True),
93+
sa.Column("claims", postgresql.JSONB(), nullable=True),
94+
sa.Column("generated_images", postgresql.JSONB(), nullable=True),
95+
sa.Column("additional_data", postgresql.JSONB(), nullable=True),
96+
sa.PrimaryKeyConstraint("id"),
97+
sa.ForeignKeyConstraint(
98+
["primary_question_id", "iteration_nr"],
99+
[
100+
"research_agent_iteration.primary_question_id",
101+
"research_agent_iteration.iteration_nr",
102+
],
103+
ondelete="CASCADE",
104+
),
105+
)
106+
107+
108+
def downgrade() -> None:
109+
# Drop tables in reverse order
110+
op.drop_table("research_agent_iteration_sub_step")
111+
op.drop_table("research_agent_iteration")
112+
113+
# Remove columns from chat_message table
114+
op.drop_column("chat_message", "research_plan")
115+
op.drop_column("chat_message", "research_type")

‎backend/alembic/versions/7ed603b64d5a_add_mcp_server_and_connection_config_.py‎

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
"""add_mcp_server_and_connection_config_models
2+
3+
Revision ID: 7ed603b64d5a
4+
Revises: b329d00a9ea6
5+
Create Date: 2025-07-28 17:35:59.900680
6+
7+
"""
8+
9+
from alembic import op
10+
import sqlalchemy as sa
11+
from sqlalchemy.dialects import postgresql
12+
from onyx.db.enums import MCPAuthenticationType
13+
14+
# revision identifiers, used by Alembic.
15+
revision = "7ed603b64d5a"
16+
down_revision = "b329d00a9ea6"
17+
branch_labels = None
18+
depends_on = None
19+
20+
21+
def upgrade() -> None:
22+
"""Create tables and columns for MCP Server support"""
23+
24+
# 1. MCP Server main table (no FK constraints yet to avoid circular refs)
25+
op.create_table(
26+
"mcp_server",
27+
sa.Column("id", sa.Integer(), primary_key=True),
28+
sa.Column("owner", sa.String(), nullable=False),
29+
sa.Column("name", sa.String(), nullable=False),
30+
sa.Column("description", sa.String(), nullable=True),
31+
sa.Column("server_url", sa.String(), nullable=False),
32+
sa.Column(
33+
"auth_type",
34+
sa.Enum(
35+
MCPAuthenticationType,
36+
name="mcp_authentication_type",
37+
native_enum=False,
38+
),
39+
nullable=False,
40+
),
41+
sa.Column("admin_connection_config_id", sa.Integer(), nullable=True),
42+
sa.Column(
43+
"created_at",
44+
sa.DateTime(timezone=True),
45+
server_default=sa.text("now()"), # type: ignore
46+
nullable=False,
47+
),
48+
sa.Column(
49+
"updated_at",
50+
sa.DateTime(timezone=True),
51+
server_default=sa.text("now()"), # type: ignore
52+
nullable=False,
53+
),
54+
)
55+
56+
# 2. MCP Connection Config table (can reference mcp_server now that it exists)
57+
op.create_table(
58+
"mcp_connection_config",
59+
sa.Column("id", sa.Integer(), primary_key=True),
60+
sa.Column("mcp_server_id", sa.Integer(), nullable=True),
61+
sa.Column("user_email", sa.String(), nullable=False, default=""),
62+
sa.Column("config", sa.LargeBinary(), nullable=False),
63+
sa.Column(
64+
"created_at",
65+
sa.DateTime(timezone=True),
66+
server_default=sa.text("now()"), # type: ignore
67+
nullable=False,
68+
),
69+
sa.Column(
70+
"updated_at",
71+
sa.DateTime(timezone=True),
72+
server_default=sa.text("now()"), # type: ignore
73+
nullable=False,
74+
),
75+
sa.ForeignKeyConstraint(
76+
["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
77+
),
78+
)
79+
80+
# Helpful indexes
81+
op.create_index(
82+
"ix_mcp_connection_config_server_user",
83+
"mcp_connection_config",
84+
["mcp_server_id", "user_email"],
85+
)
86+
op.create_index(
87+
"ix_mcp_connection_config_user_email",
88+
"mcp_connection_config",
89+
["user_email"],
90+
)
91+
92+
# 3. Add the back-references from mcp_server to connection configs
93+
op.create_foreign_key(
94+
"mcp_server_admin_config_fk",
95+
"mcp_server",
96+
"mcp_connection_config",
97+
["admin_connection_config_id"],
98+
["id"],
99+
ondelete="SET NULL",
100+
)
101+
102+
# 4. Association / access-control tables
103+
op.create_table(
104+
"mcp_server__user",
105+
sa.Column("mcp_server_id", sa.Integer(), primary_key=True),
106+
sa.Column("user_id", sa.UUID(), primary_key=True),
107+
sa.ForeignKeyConstraint(
108+
["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
109+
),
110+
sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
111+
)
112+
113+
op.create_table(
114+
"mcp_server__user_group",
115+
sa.Column("mcp_server_id", sa.Integer(), primary_key=True),
116+
sa.Column("user_group_id", sa.Integer(), primary_key=True),
117+
sa.ForeignKeyConstraint(
118+
["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
119+
),
120+
sa.ForeignKeyConstraint(["user_group_id"], ["user_group.id"]),
121+
)
122+
123+
# 5. Update existing `tool` table – allow tools to belong to an MCP server
124+
op.add_column(
125+
"tool",
126+
sa.Column("mcp_server_id", sa.Integer(), nullable=True),
127+
)
128+
# Add column for MCP tool input schema
129+
op.add_column(
130+
"tool",
131+
sa.Column("mcp_input_schema", postgresql.JSONB(), nullable=True),
132+
)
133+
op.create_foreign_key(
134+
"tool_mcp_server_fk",
135+
"tool",
136+
"mcp_server",
137+
["mcp_server_id"],
138+
["id"],
139+
ondelete="CASCADE",
140+
)
141+
142+
# 6. Update persona__tool foreign keys to cascade delete
143+
# This ensures that when a tool is deleted (including via MCP server deletion),
144+
# the corresponding persona__tool rows are also deleted
145+
op.drop_constraint(
146+
"persona__tool_tool_id_fkey", "persona__tool", type_="foreignkey"
147+
)
148+
op.drop_constraint(
149+
"persona__tool_persona_id_fkey", "persona__tool", type_="foreignkey"
150+
)
151+
152+
op.create_foreign_key(
153+
"persona__tool_persona_id_fkey",
154+
"persona__tool",
155+
"persona",
156+
["persona_id"],
157+
["id"],
158+
ondelete="CASCADE",
159+
)
160+
op.create_foreign_key(
161+
"persona__tool_tool_id_fkey",
162+
"persona__tool",
163+
"tool",
164+
["tool_id"],
165+
["id"],
166+
ondelete="CASCADE",
167+
)
168+
169+
# 7. Update research_agent_iteration_sub_step foreign key to SET NULL on delete
170+
# This ensures that when a tool is deleted, the sub_step_tool_id is set to NULL
171+
# instead of causing a foreign key constraint violation
172+
op.drop_constraint(
173+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
174+
"research_agent_iteration_sub_step",
175+
type_="foreignkey",
176+
)
177+
op.create_foreign_key(
178+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
179+
"research_agent_iteration_sub_step",
180+
"tool",
181+
["sub_step_tool_id"],
182+
["id"],
183+
ondelete="SET NULL",
184+
)
185+
186+
187+
def downgrade() -> None:
188+
"""Drop all MCP-related tables / columns"""
189+
190+
# # # 1. Drop FK & columns from tool
191+
# op.drop_constraint("tool_mcp_server_fk", "tool", type_="foreignkey")
192+
op.execute("DELETE FROM tool WHERE mcp_server_id IS NOT NULL")
193+
194+
op.drop_constraint(
195+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
196+
"research_agent_iteration_sub_step",
197+
type_="foreignkey",
198+
)
199+
op.create_foreign_key(
200+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
201+
"research_agent_iteration_sub_step",
202+
"tool",
203+
["sub_step_tool_id"],
204+
["id"],
205+
)
206+
207+
# Restore original persona__tool foreign keys (without CASCADE)
208+
op.drop_constraint(
209+
"persona__tool_persona_id_fkey", "persona__tool", type_="foreignkey"
210+
)
211+
op.drop_constraint(
212+
"persona__tool_tool_id_fkey", "persona__tool", type_="foreignkey"
213+
)
214+
215+
op.create_foreign_key(
216+
"persona__tool_persona_id_fkey",
217+
"persona__tool",
218+
"persona",
219+
["persona_id"],
220+
["id"],
221+
)
222+
op.create_foreign_key(
223+
"persona__tool_tool_id_fkey",
224+
"persona__tool",
225+
"tool",
226+
["tool_id"],
227+
["id"],
228+
)
229+
op.drop_column("tool", "mcp_input_schema")
230+
op.drop_column("tool", "mcp_server_id")
231+
232+
# 2. Drop association tables
233+
op.drop_table("mcp_server__user_group")
234+
op.drop_table("mcp_server__user")
235+
236+
# 3. Drop FK from mcp_server to connection configs
237+
op.drop_constraint("mcp_server_admin_config_fk", "mcp_server", type_="foreignkey")
238+
239+
# 4. Drop connection config indexes & table
240+
op.drop_index(
241+
"ix_mcp_connection_config_user_email", table_name="mcp_connection_config"
242+
)
243+
op.drop_index(
244+
"ix_mcp_connection_config_server_user", table_name="mcp_connection_config"
245+
)
246+
op.drop_table("mcp_connection_config")
247+
248+
# 5. Finally drop mcp_server table
249+
op.drop_table("mcp_server")

0 commit comments

Comments
 (0)