Skip to content

Commit e279b8f

Browse files
crivetimihaicrivetimihai
authored
Fix postgresql migration (#817)
* Fix postgresql migration Signed-off-by: Mihai Criveti <[email protected]> * Fix postgresql migration Signed-off-by: Mihai Criveti <[email protected]> * Fix postgresql migration ALTER TABLE resources ADD COLUMN created_by VARCHAR issue Signed-off-by: Mihai Criveti <[email protected]> * Fix additional issues wiht postgresql migration Signed-off-by: Mihai Criveti <[email protected]> --------- Signed-off-by: Mihai Criveti <[email protected]>
1 parent 1b7d828 commit e279b8f

8 files changed

+429
-170
lines changed

‎mcpgateway/alembic/versions/34492f99a0c4_add_comprehensive_metadata_to_all_.py

Lines changed: 117 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -27,61 +27,130 @@
2727

2828
def upgrade() -> None:
2929
"""Add comprehensive metadata columns to all entity tables for audit tracking."""
30+
bind = op.get_bind()
31+
inspector = sa.inspect(bind)
32+
33+
# Check if this is a fresh database without existing tables
34+
if not inspector.has_table("gateways"):
35+
print("Fresh database detected. Skipping metadata migration.")
36+
return
37+
3038
tables = ["tools", "resources", "prompts", "servers", "gateways"]
3139

40+
# Define metadata columns to add
41+
metadata_columns = [
42+
("created_by", sa.String(), True),
43+
("created_from_ip", sa.String(), True),
44+
("created_via", sa.String(), True),
45+
("created_user_agent", sa.Text(), True),
46+
("modified_by", sa.String(), True),
47+
("modified_from_ip", sa.String(), True),
48+
("modified_via", sa.String(), True),
49+
("modified_user_agent", sa.Text(), True),
50+
("import_batch_id", sa.String(), True),
51+
("federation_source", sa.String(), True),
52+
("version", sa.Integer(), False, "1"), # Not nullable, with default
53+
]
54+
55+
# Add columns to each table if they don't exist
56+
for table in tables:
57+
if inspector.has_table(table):
58+
columns = [col["name"] for col in inspector.get_columns(table)]
59+
60+
for col_name, col_type, nullable, *default in metadata_columns:
61+
if col_name not in columns:
62+
try:
63+
if default:
64+
op.add_column(table, sa.Column(col_name, col_type, nullable=nullable, server_default=default[0]))
65+
else:
66+
op.add_column(table, sa.Column(col_name, col_type, nullable=nullable))
67+
print(f"Added column {col_name} to {table}")
68+
except Exception as e:
69+
print(f"Warning: Could not add column {col_name} to {table}: {e}")
70+
71+
# Create indexes for query performance (safe B-tree indexes)
72+
# Note: modified_at column doesn't exist in schema, so we skip it
73+
index_definitions = [
74+
("created_by", ["created_by"]),
75+
("created_at", ["created_at"]),
76+
("created_via", ["created_via"]),
77+
]
78+
3279
for table in tables:
33-
# Creation metadata (nullable=True for backwards compatibility)
34-
op.add_column(table, sa.Column("created_by", sa.String(), nullable=True))
35-
op.add_column(table, sa.Column("created_from_ip", sa.String(), nullable=True))
36-
op.add_column(table, sa.Column("created_via", sa.String(), nullable=True))
37-
op.add_column(table, sa.Column("created_user_agent", sa.Text(), nullable=True))
38-
39-
# Modification metadata (nullable=True for backwards compatibility)
40-
op.add_column(table, sa.Column("modified_by", sa.String(), nullable=True))
41-
op.add_column(table, sa.Column("modified_from_ip", sa.String(), nullable=True))
42-
op.add_column(table, sa.Column("modified_via", sa.String(), nullable=True))
43-
op.add_column(table, sa.Column("modified_user_agent", sa.Text(), nullable=True))
44-
45-
# Source tracking (nullable=True for backwards compatibility)
46-
op.add_column(table, sa.Column("import_batch_id", sa.String(), nullable=True))
47-
op.add_column(table, sa.Column("federation_source", sa.String(), nullable=True))
48-
op.add_column(table, sa.Column("version", sa.Integer(), nullable=False, server_default="1"))
49-
50-
# Create indexes for query performance (PostgreSQL compatible, SQLite ignores)
51-
try:
52-
op.create_index(f"idx_{table}_created_by", table, ["created_by"])
53-
op.create_index(f"idx_{table}_created_at", table, ["created_at"])
54-
op.create_index(f"idx_{table}_modified_at", table, ["modified_at"])
55-
op.create_index(f"idx_{table}_created_via", table, ["created_via"])
56-
except Exception: # nosec B110 - database compatibility
57-
# SQLite doesn't support all index types, skip silently
58-
pass
80+
if inspector.has_table(table):
81+
try:
82+
existing_indexes = [idx["name"] for idx in inspector.get_indexes(table)]
83+
except Exception as e:
84+
print(f"Warning: Could not get indexes for {table}: {e}")
85+
continue
86+
87+
for index_suffix, columns in index_definitions:
88+
index_name = f"idx_{table}_{index_suffix}"
89+
if index_name not in existing_indexes:
90+
# Check if the column exists before creating index
91+
table_columns = [col["name"] for col in inspector.get_columns(table)]
92+
if all(col in table_columns for col in columns):
93+
try:
94+
op.create_index(index_name, table, columns)
95+
print(f"Created index {index_name}")
96+
except Exception as e:
97+
print(f"Warning: Could not create index {index_name}: {e}")
98+
else:
99+
print(f"Skipping index {index_name} - required columns {columns} not found in {table}")
59100

60101

61102
def downgrade() -> None:
62103
"""Remove comprehensive metadata columns from all entity tables."""
104+
bind = op.get_bind()
105+
inspector = sa.inspect(bind)
106+
63107
tables = ["tools", "resources", "prompts", "servers", "gateways"]
64108

109+
# Index names to drop (modified_at doesn't exist, so skip it)
110+
index_suffixes = ["created_by", "created_at", "created_via"]
111+
112+
# Drop indexes first (if they exist)
65113
for table in tables:
66-
# Drop indexes first (if they exist)
67-
try:
68-
op.drop_index(f"idx_{table}_created_by", table)
69-
op.drop_index(f"idx_{table}_created_at", table)
70-
op.drop_index(f"idx_{table}_modified_at", table)
71-
op.drop_index(f"idx_{table}_created_via", table)
72-
except Exception: # nosec B110 - database compatibility
73-
# Indexes might not exist on SQLite
74-
pass
75-
76-
# Drop metadata columns
77-
op.drop_column(table, "version")
78-
op.drop_column(table, "federation_source")
79-
op.drop_column(table, "import_batch_id")
80-
op.drop_column(table, "modified_user_agent")
81-
op.drop_column(table, "modified_via")
82-
op.drop_column(table, "modified_from_ip")
83-
op.drop_column(table, "modified_by")
84-
op.drop_column(table, "created_user_agent")
85-
op.drop_column(table, "created_via")
86-
op.drop_column(table, "created_from_ip")
87-
op.drop_column(table, "created_by")
114+
if inspector.has_table(table):
115+
try:
116+
existing_indexes = [idx["name"] for idx in inspector.get_indexes(table)]
117+
except Exception as e:
118+
print(f"Warning: Could not get indexes for {table}: {e}")
119+
continue
120+
121+
for suffix in index_suffixes:
122+
index_name = f"idx_{table}_{suffix}"
123+
if index_name in existing_indexes:
124+
try:
125+
op.drop_index(index_name, table)
126+
print(f"Dropped index {index_name}")
127+
except Exception as e:
128+
print(f"Warning: Could not drop index {index_name}: {e}")
129+
130+
# Metadata columns to drop (in reverse order for safety)
131+
metadata_columns = [
132+
"version",
133+
"federation_source",
134+
"import_batch_id",
135+
"modified_user_agent",
136+
"modified_via",
137+
"modified_from_ip",
138+
"modified_by",
139+
"created_user_agent",
140+
"created_via",
141+
"created_from_ip",
142+
"created_by",
143+
]
144+
145+
# Drop metadata columns (if they exist)
146+
for table in reversed(tables): # Reverse order for safety
147+
if inspector.has_table(table):
148+
columns = [col["name"] for col in inspector.get_columns(table)]
149+
150+
for col_name in metadata_columns:
151+
if col_name in columns:
152+
try:
153+
op.drop_column(table, col_name)
154+
print(f"Dropped column {col_name} from {table}")
155+
except Exception as e:
156+
print(f"Warning: Could not drop column {col_name} from {table}: {e}")

‎mcpgateway/alembic/versions/3b17fdc40a8d_add_passthrough_headers_to_gateways_and_.py

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,17 +27,36 @@
2727

2828
def upgrade() -> None:
2929
"""Upgrade schema."""
30-
# Create global_config table
31-
op.create_table("global_config", sa.Column("id", sa.Integer(), nullable=False), sa.Column("passthrough_headers", sa.JSON(), nullable=True), sa.PrimaryKeyConstraint("id"))
30+
bind = op.get_bind()
31+
inspector = sa.inspect(bind)
3232

33-
# Add passthrough_headers column to gateways table
34-
op.add_column("gateways", sa.Column("passthrough_headers", sa.JSON(), nullable=True))
33+
# Check if this is a fresh database without existing tables
34+
if not inspector.has_table("gateways"):
35+
print("Fresh database detected. Skipping passthrough headers migration.")
36+
return
37+
38+
# Create global_config table if it doesn't exist
39+
if not inspector.has_table("global_config"):
40+
op.create_table("global_config", sa.Column("id", sa.Integer(), nullable=False), sa.Column("passthrough_headers", sa.JSON(), nullable=True), sa.PrimaryKeyConstraint("id"))
41+
42+
# Add passthrough_headers column to gateways table if it doesn't exist
43+
if inspector.has_table("gateways"):
44+
columns = [col["name"] for col in inspector.get_columns("gateways")]
45+
if "passthrough_headers" not in columns:
46+
op.add_column("gateways", sa.Column("passthrough_headers", sa.JSON(), nullable=True))
3547

3648

3749
def downgrade() -> None:
3850
"""Downgrade schema."""
39-
# Remove passthrough_headers column from gateways table
40-
op.drop_column("gateways", "passthrough_headers")
41-
42-
# Drop global_config table
43-
op.drop_table("global_config")
51+
bind = op.get_bind()
52+
inspector = sa.inspect(bind)
53+
54+
# Remove passthrough_headers column from gateways table if it exists
55+
if inspector.has_table("gateways"):
56+
columns = [col["name"] for col in inspector.get_columns("gateways")]
57+
if "passthrough_headers" in columns:
58+
op.drop_column("gateways", "passthrough_headers")
59+
60+
# Drop global_config table if it exists
61+
if inspector.has_table("global_config"):
62+
op.drop_table("global_config")

‎mcpgateway/alembic/versions/add_a2a_agents_and_metrics.py

Lines changed: 79 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ def upgrade() -> None:
3434
existing_tables = inspector.get_table_names()
3535

3636
if "a2a_agents" not in existing_tables:
37-
# Create a2a_agents table
37+
# Create a2a_agents table with unique constraints included (SQLite compatible)
3838
op.create_table(
3939
"a2a_agents",
4040
sa.Column("id", sa.String(36), primary_key=True),
@@ -65,12 +65,10 @@ def upgrade() -> None:
6565
sa.Column("import_batch_id", sa.String()),
6666
sa.Column("federation_source", sa.String()),
6767
sa.Column("version", sa.Integer(), nullable=False, server_default="1"),
68+
sa.UniqueConstraint("name", name="uq_a2a_agents_name"),
69+
sa.UniqueConstraint("slug", name="uq_a2a_agents_slug"),
6870
)
6971

70-
# Create unique constraints
71-
op.create_unique_constraint("uq_a2a_agents_name", "a2a_agents", ["name"])
72-
op.create_unique_constraint("uq_a2a_agents_slug", "a2a_agents", ["slug"])
73-
7472
if "a2a_agent_metrics" not in existing_tables:
7573
# Create a2a_agent_metrics table
7674
op.create_table(
@@ -93,57 +91,90 @@ def upgrade() -> None:
9391
)
9492

9593
# Create indexes for performance (check if they exist first)
96-
existing_indexes = []
97-
try:
98-
existing_indexes = [idx["name"] for idx in inspector.get_indexes("a2a_agents")]
99-
except Exception:
100-
pass
101-
102-
if "idx_a2a_agents_enabled" not in existing_indexes:
94+
# Only create indexes if tables were actually created
95+
if "a2a_agents" in existing_tables:
10396
try:
104-
op.create_index("idx_a2a_agents_enabled", "a2a_agents", ["enabled"])
97+
existing_indexes = [idx["name"] for idx in inspector.get_indexes("a2a_agents")]
10598
except Exception:
106-
pass
99+
existing_indexes = []
100+
101+
if "idx_a2a_agents_enabled" not in existing_indexes:
102+
try:
103+
op.create_index("idx_a2a_agents_enabled", "a2a_agents", ["enabled"])
104+
except Exception as e:
105+
print(f"Warning: Could not create index idx_a2a_agents_enabled: {e}")
106+
107+
if "idx_a2a_agents_agent_type" not in existing_indexes:
108+
try:
109+
op.create_index("idx_a2a_agents_agent_type", "a2a_agents", ["agent_type"])
110+
except Exception as e:
111+
print(f"Warning: Could not create index idx_a2a_agents_agent_type: {e}")
107112

108-
if "idx_a2a_agents_agent_type" not in existing_indexes:
113+
# Create B-tree index for tags (safer than GIN, works on both PostgreSQL and SQLite)
114+
if "idx_a2a_agents_tags" not in existing_indexes:
115+
try:
116+
op.create_index("idx_a2a_agents_tags", "a2a_agents", ["tags"])
117+
except Exception as e:
118+
print(f"Warning: Could not create index idx_a2a_agents_tags: {e}")
119+
120+
# Metrics table indexes
121+
if "a2a_agent_metrics" in existing_tables:
109122
try:
110-
op.create_index("idx_a2a_agents_agent_type", "a2a_agents", ["agent_type"])
123+
existing_metrics_indexes = [idx["name"] for idx in inspector.get_indexes("a2a_agent_metrics")]
111124
except Exception:
112-
pass
125+
existing_metrics_indexes = []
113126

114-
# Metrics table indexes
115-
try:
116-
existing_indexes = [idx["name"] for idx in inspector.get_indexes("a2a_agent_metrics")]
117-
if "idx_a2a_agent_metrics_agent_id" not in existing_indexes:
118-
op.create_index("idx_a2a_agent_metrics_agent_id", "a2a_agent_metrics", ["a2a_agent_id"])
119-
if "idx_a2a_agent_metrics_timestamp" not in existing_indexes:
120-
op.create_index("idx_a2a_agent_metrics_timestamp", "a2a_agent_metrics", ["timestamp"])
121-
except Exception:
122-
pass
123-
124-
# Create GIN indexes for tags on PostgreSQL (ignored on SQLite)
125-
try:
126-
if "idx_a2a_agents_tags" not in existing_indexes:
127-
op.create_index("idx_a2a_agents_tags", "a2a_agents", ["tags"], postgresql_using="gin")
128-
except Exception: # nosec B110 - database compatibility
129-
pass # SQLite doesn't support GIN indexes
127+
if "idx_a2a_agent_metrics_agent_id" not in existing_metrics_indexes:
128+
try:
129+
op.create_index("idx_a2a_agent_metrics_agent_id", "a2a_agent_metrics", ["a2a_agent_id"])
130+
except Exception as e:
131+
print(f"Warning: Could not create index idx_a2a_agent_metrics_agent_id: {e}")
132+
133+
if "idx_a2a_agent_metrics_timestamp" not in existing_metrics_indexes:
134+
try:
135+
op.create_index("idx_a2a_agent_metrics_timestamp", "a2a_agent_metrics", ["timestamp"])
136+
except Exception as e:
137+
print(f"Warning: Could not create index idx_a2a_agent_metrics_timestamp: {e}")
130138

131139

132140
def downgrade() -> None:
133141
"""Reverse the A2A agents and metrics tables."""
142+
# Check if tables exist before trying to drop indexes/tables
143+
conn = op.get_bind()
144+
inspector = sa.inspect(conn)
145+
existing_tables = inspector.get_table_names()
134146

135-
# Drop indexes first
136-
try:
137-
op.drop_index("idx_a2a_agents_tags", "a2a_agents")
138-
except Exception: # nosec B110 - database compatibility
139-
pass
140-
141-
op.drop_index("idx_a2a_agent_metrics_timestamp", "a2a_agent_metrics")
142-
op.drop_index("idx_a2a_agent_metrics_agent_id", "a2a_agent_metrics")
143-
op.drop_index("idx_a2a_agents_agent_type", "a2a_agents")
144-
op.drop_index("idx_a2a_agents_enabled", "a2a_agents")
145-
146-
# Drop tables
147-
op.drop_table("server_a2a_association")
148-
op.drop_table("a2a_agent_metrics")
149-
op.drop_table("a2a_agents")
147+
# Drop indexes first (if they exist)
148+
if "a2a_agents" in existing_tables:
149+
try:
150+
existing_indexes = [idx["name"] for idx in inspector.get_indexes("a2a_agents")]
151+
152+
for index_name in ["idx_a2a_agents_tags", "idx_a2a_agents_agent_type", "idx_a2a_agents_enabled"]:
153+
if index_name in existing_indexes:
154+
try:
155+
op.drop_index(index_name, "a2a_agents")
156+
except Exception as e:
157+
print(f"Warning: Could not drop index {index_name}: {e}")
158+
except Exception as e:
159+
print(f"Warning: Could not get indexes for a2a_agents: {e}")
160+
161+
if "a2a_agent_metrics" in existing_tables:
162+
try:
163+
existing_metrics_indexes = [idx["name"] for idx in inspector.get_indexes("a2a_agent_metrics")]
164+
165+
for index_name in ["idx_a2a_agent_metrics_timestamp", "idx_a2a_agent_metrics_agent_id"]:
166+
if index_name in existing_metrics_indexes:
167+
try:
168+
op.drop_index(index_name, "a2a_agent_metrics")
169+
except Exception as e:
170+
print(f"Warning: Could not drop index {index_name}: {e}")
171+
except Exception as e:
172+
print(f"Warning: Could not get indexes for a2a_agent_metrics: {e}")
173+
174+
# Drop tables (if they exist)
175+
for table_name in ["server_a2a_association", "a2a_agent_metrics", "a2a_agents"]:
176+
if table_name in existing_tables:
177+
try:
178+
op.drop_table(table_name)
179+
except Exception as e:
180+
print(f"Warning: Could not drop table {table_name}: {e}")

0 commit comments

Comments
 (0)