fixed missing primary key import for LazyTableReference

[sinisaos]

Related to #865. Adds missing primary key to extra_imports when creating a migration with a ForeignKey that uses a LazyTableReference.

[dantownsend]

Thanks for this - I tested it, and it works.

I think the problem goes even deeper though.

Let's say you have this (note the UUID primary key):

class TableA(Table):
    table_b = ForeignKey(LazyTableReference('TableB', module_path=__module__))

class TableB(Table):
    id = UUID(primary_key=True)

If you create the migration, it now has the import for UUID, but it's missing the import for UUID4 (the default value for UUID).

The problem is, we use SerialisedTableType, which itself calls serialise_params. We need to access the extra_imports from that serialise_params call.

Sorry, it sounds very confusing.

But basically, I think the ideal solution would be this:

serialised_table_type = SerialisedTableType(table_type=table_type)
extra_definitions.append(serialised_table_type)
extra_imports.extend(serialised_table_type.extra_imports)

So we need to refactor SerialisedTableType slightly so can access its extra imports.

Does that make sense?

[sinisaos]

In my case everything works fine because I already tried this PR with custom primary key. UUID4() is imported along with UUID with tables like this.

class TableA(Table):
    id = UUID(primary_key=True)
    table_b = ForeignKey(LazyTableReference("TableB", module_path=__module__))

class TableB(Table):
    id = UUID(primary_key=True)

Migration file

from piccolo.apps.migrations.auto.migration_manager import MigrationManager
from piccolo.columns.base import OnDelete
from piccolo.columns.base import OnUpdate
from piccolo.columns.column_types import ForeignKey
from piccolo.columns.column_types import UUID # <- correct import
from piccolo.columns.defaults.uuid import UUID4 # <- correct import
from piccolo.columns.indexes import IndexMethod
from piccolo.table import Table


class TableB(Table, tablename="table_b", schema=None):
    id = UUID( # <- correct
        default=UUID4(), # <- correct
        null=False,
        primary_key=True,
        unique=False,
        index=False,
        index_method=IndexMethod.btree,
        choices=None,
        db_column_name=None,
        secret=False,
    )


ID = "2025-07-31T09:41:33:395378"
VERSION = "1.28.0"
DESCRIPTION = ""


async def forwards():
    manager = MigrationManager(
        migration_id=ID, app_name="home", description=DESCRIPTION
    )

    manager.add_table(
        class_name="TableB", tablename="table_b", schema=None, columns=None
    )

    manager.add_table(
        class_name="TableA", tablename="table_a", schema=None, columns=None
    )

    manager.add_column(
        table_class_name="TableB",
        tablename="table_b",
        column_name="id",
        db_column_name="id",
        column_class_name="UUID",
        column_class=UUID, # <- correct
        params={
            "default": UUID4(), # <- correct
            "null": False,
            "primary_key": True,
            "unique": False,
            "index": False,
            "index_method": IndexMethod.btree,
            "choices": None,
            "db_column_name": None,
            "secret": False,
        },
        schema=None,
    )

    manager.add_column(
        table_class_name="TableA",
        tablename="table_a",
        column_name="id",
        db_column_name="id",
        column_class_name="UUID",
        column_class=UUID, # <- correct
        params={
            "default": UUID4(), # <- correct
            "null": False,
            "primary_key": True,
            "unique": False,
            "index": False,
            "index_method": IndexMethod.btree,
            "choices": None,
            "db_column_name": None,
            "secret": False,
        },
        schema=None,
    )

    manager.add_column(
        table_class_name="TableA",
        tablename="table_a",
        column_name="table_b",
        db_column_name="table_b",
        column_class_name="ForeignKey",
        column_class=ForeignKey,
        params={
            "references": TableB,
            "on_delete": OnDelete.cascade,
            "on_update": OnUpdate.cascade,
            "target_column": None,
            "null": True,
            "primary_key": False,
            "unique": False,
            "index": False,
            "index_method": IndexMethod.btree,
            "choices": None,
            "db_column_name": None,
            "secret": False,
        },
        schema=None,
    )

    return manager

If you think the SerialisedTableType should be changed, feel free to implement it in this PR.

[Skelmis]

Tested with serial and uuid and it works. Will merge as is given the length of time since last comment, we can raise another PR in future if required but I'd prefer to just get the functionality in at a working base level