From a0fde50307fb10761f384c345601e146c045bed4 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Thu, 10 Apr 2025 17:54:33 +0200 Subject: [PATCH 1/4] Update alembic version to "^1.8.1" in pyproject.toml --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b8bf06c5f3..b2303934af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ include = ["src/zenml", "*.txt", "*.sh", "*.md"] zenml = "zenml.cli.cli:cli" [tool.poetry.dependencies] -alembic = { version = "~1.8.1" } +alembic = { version = "^1.8.1" } bcrypt = { version = "4.0.1" } click = "^8.0.1,<8.1.8" cloudpickle = ">=2.0.0,<3" From ce6605ac0f069ad7caa5726e15c5f9d64e02901a Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Thu, 10 Apr 2025 22:13:07 +0200 Subject: [PATCH 2/4] Update typing and schema imports in alembic.py --- src/zenml/zen_stores/migrations/alembic.py | 71 ++++++++++++++++------ 1 file changed, 54 insertions(+), 17 deletions(-) diff --git a/src/zenml/zen_stores/migrations/alembic.py b/src/zenml/zen_stores/migrations/alembic.py index 65430cafc6..f72499805b 100644 --- a/src/zenml/zen_stores/migrations/alembic.py +++ b/src/zenml/zen_stores/migrations/alembic.py @@ -19,7 +19,16 @@ """ from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Sequence, Union +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + Optional, + Sequence, + Union, +) from alembic.config import Config from alembic.runtime.environment import EnvironmentContext @@ -28,7 +37,7 @@ from sqlalchemy import Column, String from sqlalchemy.engine import Engine from sqlalchemy.orm import declarative_base -from sqlalchemy.sql.schema import MetaData +from sqlalchemy.sql.schema import MetaData, SchemaItem from sqlmodel import SQLModel from zenml.zen_stores import schemas @@ -39,7 +48,18 @@ def include_object( - object: Any, name: str, type_: str, *args: Any, **kwargs: Any + object: SchemaItem, + name: Optional[str], + type_: Literal[ + "schema", + "table", + "column", + "index", + "unique_constraint", + "foreign_key_constraint", + ], + reflected: bool, + compare_to: Optional[SchemaItem], ) -> bool: """Function used to exclude tables from the migration scripts. @@ -47,8 +67,8 @@ def include_object( object: The schema item object to check. name: The name of the object to check. type_: The type of the object to check. - *args: Additional arguments. - **kwargs: Additional keyword arguments. + reflected: Whether this object is being reflected. + compare_to: The object being compared against. Returns: True if the object should be included, False otherwise. @@ -180,9 +200,14 @@ def current_revisions(self) -> List[str]: def do_get_current_rev(rev: _RevIdType, context: Any) -> List[Any]: nonlocal current_revisions - for r in self.script_directory.get_all_current( - rev # type:ignore [arg-type] - ): + # Always convert to tuple of strings as get_all_current expects this type + if isinstance(rev, str): + # Explicitly annotate with variable-length tuple + rev_tuple: Sequence[str] = (rev,) + else: + rev_tuple = tuple(str(r) for r in rev) + + for r in self.script_directory.get_all_current(rev_tuple): # type: ignore[arg-type] if r is None: continue current_revisions.append(r.revision) @@ -200,7 +225,11 @@ def stamp(self, revision: str) -> None: """ def do_stamp(rev: _RevIdType, context: Any) -> List[Any]: - return self.script_directory._stamp_revs(revision, rev) + if isinstance(rev, str): + return self.script_directory._stamp_revs(revision, rev) + else: + rev_tuple = tuple(str(r) for r in rev) + return self.script_directory._stamp_revs(revision, rev_tuple) self.run_migrations(do_stamp) @@ -212,10 +241,14 @@ def upgrade(self, revision: str = "heads") -> None: """ def do_upgrade(rev: _RevIdType, context: Any) -> List[Any]: - return self.script_directory._upgrade_revs( - revision, - rev, # type:ignore [arg-type] - ) + if isinstance(rev, str): + return self.script_directory._upgrade_revs(revision, rev) + else: + if rev: + return self.script_directory._upgrade_revs( + revision, str(rev[0]) + ) + return [] self.run_migrations(do_upgrade) @@ -227,9 +260,13 @@ def downgrade(self, revision: str) -> None: """ def do_downgrade(rev: _RevIdType, context: Any) -> List[Any]: - return self.script_directory._downgrade_revs( - revision, - rev, # type:ignore [arg-type] - ) + if isinstance(rev, str): + return self.script_directory._downgrade_revs(revision, rev) + else: + if rev: + return self.script_directory._downgrade_revs( + revision, str(rev[0]) + ) + return self.script_directory._downgrade_revs(revision, None) self.run_migrations(do_downgrade) From 411987c12c0a53c0c3a870997705aa25f8f44a41 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Thu, 10 Apr 2025 22:16:27 +0200 Subject: [PATCH 3/4] Update Alembic imports for backward compatibility --- src/zenml/zen_stores/migrations/alembic.py | 35 +++++++++++----------- 1 file changed, 17 insertions(+), 18 deletions(-) diff --git a/src/zenml/zen_stores/migrations/alembic.py b/src/zenml/zen_stores/migrations/alembic.py index f72499805b..a46e6481dd 100644 --- a/src/zenml/zen_stores/migrations/alembic.py +++ b/src/zenml/zen_stores/migrations/alembic.py @@ -24,7 +24,6 @@ Callable, Dict, List, - Literal, Optional, Sequence, Union, @@ -37,7 +36,7 @@ from sqlalchemy import Column, String from sqlalchemy.engine import Engine from sqlalchemy.orm import declarative_base -from sqlalchemy.sql.schema import MetaData, SchemaItem +from sqlalchemy.sql.schema import MetaData from sqlmodel import SQLModel from zenml.zen_stores import schemas @@ -48,18 +47,11 @@ def include_object( - object: SchemaItem, + object: Any, # Use Any for backward compatibility name: Optional[str], - type_: Literal[ - "schema", - "table", - "column", - "index", - "unique_constraint", - "foreign_key_constraint", - ], - reflected: bool, - compare_to: Optional[SchemaItem], + type_: str, # Use string instead of Literal for backward compatibility + reflected: bool = False, + compare_to: Optional[Any] = None, ) -> bool: """Function used to exclude tables from the migration scripts. @@ -155,6 +147,7 @@ def run_migrations( fn_context_args["fn"] = fn with self.engine.connect() as connection: + # Configure the context with our metadata self.environment_context.configure( connection=connection, target_metadata=self.metadata, @@ -200,14 +193,15 @@ def current_revisions(self) -> List[str]: def do_get_current_rev(rev: _RevIdType, context: Any) -> List[Any]: nonlocal current_revisions - # Always convert to tuple of strings as get_all_current expects this type + # Handle rev parameter in a way that's compatible with different alembic versions + rev_input: Any if isinstance(rev, str): - # Explicitly annotate with variable-length tuple - rev_tuple: Sequence[str] = (rev,) + rev_input = rev else: - rev_tuple = tuple(str(r) for r in rev) + rev_input = tuple(str(r) for r in rev) - for r in self.script_directory.get_all_current(rev_tuple): # type: ignore[arg-type] + # Get current revision(s) + for r in self.script_directory.get_all_current(rev_input): if r is None: continue current_revisions.append(r.revision) @@ -225,9 +219,11 @@ def stamp(self, revision: str) -> None: """ def do_stamp(rev: _RevIdType, context: Any) -> List[Any]: + # Handle rev parameter in a way that's compatible with different alembic versions if isinstance(rev, str): return self.script_directory._stamp_revs(revision, rev) else: + # Convert to tuple for compatibility rev_tuple = tuple(str(r) for r in rev) return self.script_directory._stamp_revs(revision, rev_tuple) @@ -241,10 +237,12 @@ def upgrade(self, revision: str = "heads") -> None: """ def do_upgrade(rev: _RevIdType, context: Any) -> List[Any]: + # Handle rev parameter in a way that's compatible with different alembic versions if isinstance(rev, str): return self.script_directory._upgrade_revs(revision, rev) else: if rev: + # Use first element or revs for compatibility return self.script_directory._upgrade_revs( revision, str(rev[0]) ) @@ -260,6 +258,7 @@ def downgrade(self, revision: str) -> None: """ def do_downgrade(rev: _RevIdType, context: Any) -> List[Any]: + # Handle rev parameter in a way that's compatible with different alembic versions if isinstance(rev, str): return self.script_directory._downgrade_revs(revision, rev) else: From 99beb2f4437da6e72400e823a2176db68cd83ed7 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Mon, 14 Apr 2025 11:06:39 +0200 Subject: [PATCH 4/4] Update alembic version range in dependencies --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b2303934af..e8ff688111 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ include = ["src/zenml", "*.txt", "*.sh", "*.md"] zenml = "zenml.cli.cli:cli" [tool.poetry.dependencies] -alembic = { version = "^1.8.1" } +alembic = { version = ">=1.8.1,<=1.15.2" } bcrypt = { version = "4.0.1" } click = "^8.0.1,<8.1.8" cloudpickle = ">=2.0.0,<3"