Skip to content

Commit

Permalink
Merge pull request #2054 from Agenta-AI/feature/age-761-bug-downgrade…
Browse files Browse the repository at this point in the history 
…-migration-fails-with-sqlalchemyexccompileerror

[Refactor]: Manually Name Database Table Constraints to Resolve sqlalchemy.exc.CompileError
  • Loading branch information
@aakrem
aakrem committed Sep 27, 2024
2 parents 239a655 + 22b95ad commit da443bb
Show file tree
Hide file tree
Showing 2 changed files with 174 additions and 43 deletions.
24 changes: 21 additions & 3 deletions ...ckend/migrations/postgres/versions/5c29a64204f4_added_modified_by_id_column_to_apps_db.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
Create Date: 2024-08-25 17:56:11.732929
"""

import os
from typing import Sequence, Union

from alembic import op
Expand All @@ -18,16 +20,32 @@
depends_on: Union[str, Sequence[str], None] = None


def unique_constraint_exists(engine: sa.Engine, table_name: str, constraint_name: str):
with engine.connect() as conn:
result = conn.execute(
sa.text(
f"""
SELECT conname FROM pg_constraint
WHERE conname = '{constraint_name}' AND conrelid = '{table_name}'::regclass;
"""
)
)
return result.fetchone() is not None


def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.add_column("app_db", sa.Column("modified_by_id", sa.UUID(), nullable=True))
op.create_foreign_key(None, "app_db", "users", ["modified_by_id"], ["id"])
op.create_foreign_key(
"app_db_modified_by_id_fkey", "app_db", "users", ["modified_by_id"], ["id"]
)
# ### end Alembic commands ###


def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint(None, "app_db", type_="foreignkey")
op.drop_constraint(None, "app_db", type_="unique")
engine = sa.create_engine(os.getenv("POSTGRES_URI")) # type: ignore
if not unique_constraint_exists(engine, "app_db", "app_db_modified_by_id_fkey"):
op.drop_constraint("app_db_modified_by_id_fkey", "app_db", type_="foreignkey")
op.drop_column("app_db", "modified_by_id")
# ### end Alembic commands ###
193 changes: 153 additions & 40 deletions agenta-backend/agenta_backend/migrations/postgres/versions/b80c708c21bb_initial_migration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
Create Date: 2024-07-11 13:20:31.705894
"""

import os
from typing import Sequence, Union

from alembic import op
Expand All @@ -23,6 +25,19 @@
depends_on: Union[str, Sequence[str], None] = None


def unique_constraint_exists(engine: sa.Engine, table_name: str, constraint_name: str):
with engine.connect() as conn:
result = conn.execute(
sa.text(
f"""
SELECT conname FROM pg_constraint
WHERE conname = '{constraint_name}' AND conrelid = '{table_name}'::regclass;
"""
)
)
return result.fetchone() is not None


def first_time_user_from_agenta_v019_upwards_upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.create_table(
Expand Down Expand Up @@ -475,26 +490,104 @@ def first_time_user_from_agenta_v019_upwards_upgrade() -> None:

def returning_user_from_agenta_v018_downwards_upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.create_unique_constraint(None, "app_db", ["id"])
op.create_unique_constraint(None, "app_variant_revisions", ["id"])
op.create_unique_constraint(None, "app_variants", ["id"])
op.create_unique_constraint(None, "bases", ["id"])
op.create_unique_constraint(None, "deployments", ["id"])
op.create_unique_constraint(None, "docker_images", ["id"])
op.create_unique_constraint(None, "environments", ["id"])
op.create_unique_constraint(None, "environments_revisions", ["id"])
op.create_unique_constraint(None, "evaluation_aggregated_results", ["id"])
op.create_unique_constraint(None, "evaluation_scenario_results", ["id"])
op.create_unique_constraint(None, "evaluation_scenarios", ["id"])
op.create_unique_constraint(None, "evaluations", ["id"])
op.create_unique_constraint(None, "evaluators_configs", ["id"])
op.create_unique_constraint(None, "human_evaluation_variants", ["id"])
op.create_unique_constraint(None, "human_evaluations", ["id"])
op.create_unique_constraint(None, "human_evaluations_scenarios", ["id"])
op.create_unique_constraint(None, "ids_mapping", ["id"])
op.create_unique_constraint(None, "templates", ["id"])
op.create_unique_constraint(None, "testsets", ["id"])
op.create_unique_constraint(None, "users", ["id"])
engine = sa.create_engine(os.getenv("POSTGRES_URI")) # type: ignore
if not unique_constraint_exists(engine, "app_db", "app_db_pkey"):
op.create_unique_constraint("app_db_pkey", "app_db", ["id"])

if not unique_constraint_exists(
engine, "app_variant_revisions", "app_variant_revisions_pkey"
):
op.create_unique_constraint(
"app_variant_revisions_pkey", "app_variant_revisions", ["id"]
)

if not unique_constraint_exists(engine, "app_variants", "app_variants_pkey"):
op.create_unique_constraint("app_variants_pkey", "app_variants", ["id"])

if not unique_constraint_exists(engine, "bases", "bases_pkey"):
op.create_unique_constraint("bases_pkey", "bases", ["id"])

if not unique_constraint_exists(engine, "deployments", "deployments_pkey"):
op.create_unique_constraint("deployments_pkey", "deployments", ["id"])

if not unique_constraint_exists(engine, "docker_images", "docker_images_pkey"):
op.create_unique_constraint("docker_images_pkey", "docker_images", ["id"])

if not unique_constraint_exists(engine, "environments", "environments_pkey"):
op.create_unique_constraint("environments_pkey", "environments", ["id"])

if not unique_constraint_exists(
engine, "environments_revisions", "environments_revisions_pkey"
):
op.create_unique_constraint(
"environments_revisions_pkey", "environments_revisions", ["id"]
)

if not unique_constraint_exists(
engine, "evaluation_aggregated_results", "evaluation_aggregated_results_pkey"
):
op.create_unique_constraint(
"evaluation_aggregated_results_pkey",
"evaluation_aggregated_results",
["id"],
)

if not unique_constraint_exists(
engine, "evaluation_scenario_results", "evaluation_scenario_results_pkey"
):
op.create_unique_constraint(
"evaluation_scenario_results_pkey", "evaluation_scenario_results", ["id"]
)

if not unique_constraint_exists(
engine, "evaluation_scenarios", "evaluation_scenarios_pkey"
):
op.create_unique_constraint(
"evaluation_scenarios_pkey", "evaluation_scenarios", ["id"]
)

if not unique_constraint_exists(engine, "evaluations", "evaluations_pkey"):
op.create_unique_constraint("evaluations_pkey", "evaluations", ["id"])

if not unique_constraint_exists(
engine, "evaluators_configs", "evaluators_configs_pkey"
):
op.create_unique_constraint(
"evaluators_configs_pkey", "evaluators_configs", ["id"]
)

if not unique_constraint_exists(
engine, "human_evaluation_variants", "human_evaluation_variants_pkey"
):
op.create_unique_constraint(
"human_evaluation_variants_pkey", "human_evaluation_variants", ["id"]
)

if not unique_constraint_exists(
engine, "human_evaluations", "human_evaluations_pkey"
):
op.create_unique_constraint(
"human_evaluations_pkey", "human_evaluations", ["id"]
)

if not unique_constraint_exists(
engine, "human_evaluations_scenarios", "human_evaluations_scenarios_pkey"
):
op.create_unique_constraint(
"human_evaluations_scenarios_pkey", "human_evaluations_scenarios", ["id"]
)

if not unique_constraint_exists(engine, "ids_mapping", "ids_mapping_pkey"):
op.create_unique_constraint("ids_mapping_pkey", "ids_mapping", ["id"])

if not unique_constraint_exists(engine, "templates", "templates_pkey"):
op.create_unique_constraint("templates_pkey", "templates", ["id"])

if not unique_constraint_exists(engine, "testsets", "testsets_pkey"):
op.create_unique_constraint("testsets_pkey", "testsets", ["id"])

if not unique_constraint_exists(engine, "users", "users_pkey"):
op.create_unique_constraint("users_pkey", "users", ["id"])
# ### end Alembic commands ###


Expand Down Expand Up @@ -528,26 +621,46 @@ def first_time_user_from_agenta_v019_upwards_downgrade() -> None:

def returning_user_from_agenta_v018_downwards_downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint(None, "users", type_="unique")
op.drop_constraint(None, "testsets", type_="unique")
op.drop_constraint(None, "templates", type_="unique")
op.drop_constraint(None, "ids_mapping", type_="unique")
op.drop_constraint(None, "human_evaluations_scenarios", type_="unique")
op.drop_constraint(None, "human_evaluations", type_="unique")
op.drop_constraint(None, "human_evaluation_variants", type_="unique")
op.drop_constraint(None, "evaluators_configs", type_="unique")
op.drop_constraint(None, "evaluations", type_="unique")
op.drop_constraint(None, "evaluation_scenarios", type_="unique")
op.drop_constraint(None, "evaluation_scenario_results", type_="unique")
op.drop_constraint(None, "evaluation_aggregated_results", type_="unique")
op.drop_constraint(None, "environments_revisions", type_="unique")
op.drop_constraint(None, "environments", type_="unique")
op.drop_constraint(None, "docker_images", type_="unique")
op.drop_constraint(None, "deployments", type_="unique")
op.drop_constraint(None, "bases", type_="unique")
op.drop_constraint(None, "app_variants", type_="unique")
op.drop_constraint(None, "app_variant_revisions", type_="unique")
op.drop_constraint(None, "app_db", type_="unique")
op.drop_constraint("users_pkey", "users", type_="unique")
op.drop_constraint("testsets_pkey", "testsets", type_="unique")
op.drop_constraint("templates_pkey", "templates", type_="unique")
op.drop_constraint("ids_mapping_pkey", "ids_mapping", type_="unique")
op.drop_constraint(
"human_evaluations_scenarios_pkey",
"human_evaluations_scenarios",
type_="unique",
)
op.drop_constraint("human_evaluations_pkey", "human_evaluations", type_="unique")
op.drop_constraint(
"human_evaluation_variants_pkey", "human_evaluation_variants", type_="unique"
)
op.drop_constraint("evaluators_configs_pkey", "evaluators_configs", type_="unique")
op.drop_constraint("evaluations_pkey", "evaluations", type_="unique")
op.drop_constraint(
"evaluation_scenarios_pkey", "evaluation_scenarios", type_="unique"
)
op.drop_constraint(
"evaluation_scenario_results_pkey",
"evaluation_scenario_results",
type_="unique",
)
op.drop_constraint(
"evaluation_aggregated_results_pkey",
"evaluation_aggregated_results",
type_="unique",
)
op.drop_constraint(
"environments_revisions_pkey", "environments_revisions", type_="unique"
)
op.drop_constraint("environments_pkey", "environments", type_="unique")
op.drop_constraint("docker_images_pkey", "docker_images", type_="unique")
op.drop_constraint("deployments_pkey", "deployments", type_="unique")
op.drop_constraint("bases_key", "bases", type_="unique")
op.drop_constraint("app_variants_pkey", "app_variants", type_="unique")
op.drop_constraint(
"app_variant_revisions_pkey", "app_variant_revisions", type_="unique"
)
op.drop_constraint("app_db_pkey", "app_db", type_="unique")
# ### end Alembic commands ###


Expand Down

0 comments on commit da443bb

Please sign in to comment.