Alembic migration with postgres multi schema not working

[biyani701]

Describe the bug
When running alembic migration with postgres multi schema, it is not working.

Expected behavior
migration should be able to identify tables that need to be created in public schema and those that have been added in other schema say test

To Reproduce
Create 2 dbmodels, one is public and one with no schema defined.
class User(Base, TableName):
ACCOUNT_TYPES = [
('atlassian', 'Atlassian'),
('app', 'App'),
('customer', 'Customer'),
]

accountId: Mapped[str] = mapped_column(
    TEXT, primary_key=True, nullable=False,
    unique=True
)
accountType: Mapped[ChoiceType] = mapped_column(
    ChoiceType(ACCOUNT_TYPES), nullable=False
)

emailAddress: Mapped[EmailType] = mapped_column(
    EmailType, unique=True, nullable=True
)
displayName: Mapped[TEXT] = mapped_column(
    TEXT, nullable=False
)

table_args = (
{'schema': 'public'}
)

class Issue(Base):
use_snake_case = True
# Drop column list -
# issue_hierarchy_level
id = mapped_column(
BIGINT, nullable=False, primary_key=True,
# unique=True
)
key = mapped_column(TEXT, nullable=False, index=True)
parent_id = mapped_column(
BIGINT,
# ForeignKey(column="issue.id", deferrable=True, initially='DEFERRED'),
nullable=True, index=True,
)
parent_key = mapped_column(
String,
# ForeignKey(column="issue.key", deferrable=True, initially='DEFERRED'),
nullable=True, index=True,
)
summary = mapped_column(String, nullable=False)
table_args = (
UniqueConstraint("id", name="uq_issue_id"),
UniqueConstraint("key", name="uq_issue_key"),
ForeignKeyConstraint([parent_id], [id], name="fk_issue_parent_id"),
ForeignKeyConstraint([parent_key], [key], name="fk_issue_parent_key"),
)

in env.py
target_schema = context.get_x_argument(as_dictionary=True).get("tenant")

alembic revision -m "Initial version" --autogenerate
alembic -x tenant=test1 upgrade head

make change to the model and run the above 2 steps 2nd time, it fails with error table already exists for the tables created in public.

Error

# Copy error here. Please include the full stack trace.

Versions.

• OS: Windows 11
• Python: python 3.11
• Alembic: 1.13.2
• SQLAlchemy: 2.0.31
• Database: Postgres
• DBAPI: postgresql+psycopg2

Additional context
The issue is how to handle multi schema for postgres

Have a nice day!

[CaselIT]

Hi,

Have you tried looking at this section of the docs? https://alembic.sqlalchemy.org/en/latest/autogenerate.html#autogenerate-include-hooks

[biyani701]

I have tried. Below are my include_object, include_name, process_revision_directives( which i didn't understand much :-()
I wanted version_table_schema to be generated in each schema so gave it the schema name supplied with -x directirve.

I have referred to this https://alembic.sqlalchemy.org/en/latest/cookbook.html#rudimental-schema-level-multi-tenancy-for-postgresql-databases as i am using schema_translate_map

Below are the commands i am running:

alembic revision -m "Initial version" --autogenerate
alembic -x tenant=plat upgrade head

Now when i do alembic -x tenant=plat check, i get below exception also.

Pls suggest.

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Scripts\alembic.exe\__main__.py", line 7, in <module>
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Lib\site-packages\alembic\config.py", line 636, in main
    CommandLine(prog=prog).main(argv=argv)
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Lib\site-packages\alembic\config.py", line 626, in main
    self.run_cmd(cfg, options)
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Lib\site-packages\alembic\config.py", line 603, in run_cmd
    fn(
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Lib\site-packages\alembic\command.py", line 297, in check
    diffs.extend(upgrade_ops.as_diffs())
                 ^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Lib\site-packages\alembic\operations\ops.py", line 2572, in as_diffs
    return list(OpContainer._ops_as_diffs(self))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Lib\site-packages\alembic\operations\ops.py", line 2582, in _ops_as_diffs
    yield op.to_diff_tuple()
          ^^^^^^^^^^^^^^^^^^
  File "C:\Users\DELL\PycharmProjects\jiradashboard\venv\Lib\site-packages\alembic\operations\ops.py", line 91, in to_diff_tuple
    raise NotImplementedError
NotImplementedError

def include_object(
        object: SchemaItem, name,
        type_: Union[
            Literal["schema"], Literal["table"], Literal["column"], Literal["index"], Literal["unique_constraint"],
            Literal["foreign_key_constraint"]],
        reflected: bool, compare_to: Optional[SchemaItem]
):
    if type_ == "table" and name == 'alembic_version':
        return False
    else:
        print(f"reflected {reflected}: {type_} {name}")
        return True

def include_name(name: str | None, type_: str, parent_names: dict[str, str | None]) -> bool:
    print("inside include name")
    print(type_, name)
    print(target_metadata.schema)
    print(type(parent_names), parent_names)

    if type_ == "schema":
        return name in [None, 'plat']
    elif type_ == "table":
        return (
                parent_names["schema_qualified_table_name"] in
                target_metadata.tables
        )
    else:
        return True

def process_revision_directives(
        _context: MigrationContext,
        revision: str | Iterable[str | None] | Iterable[str],
        directives: list[MigrationScript],
) -> None:
    script = directives[0]
    upgrade_ops = script.upgrade_ops.ops
    print("process_revision_directives called")
    for directive in directives:
        if isinstance(directive, operations.ops.CreateTableOp):
            print(f"directive scheama = {directive.schema}")
            if directive.schema is None:
                directive.schema = get_target_schema()
            elif directive.schema == 'public':
                # Handle special cases for 'public' schema if needed
                directive.schema = get_target_schema()  # or leave it as is based on your logic
    for op in upgrade_ops:
        print(op.info)

context.configure(
            connection=connection,
            target_metadata=target_metadata,
            include_schemas=False,
            version_table_schema=target_schema,
            render_item=render_item,
            include_object=include_object,
            process_revision_directives=process_revision_directives,
            include_name=include_name,
            # render_as_batch=True,
        )

        with context.begin_transaction():
            context.run_migrations()