[BUG] Can't Index a Link

[dedeswim]

Describe the bug
I would like to have a Link inside of a document, but I also want it to be unique, for which (in my understanding), I need Indexed. However, if I try that, I get an exception as below.

To Reproduce

from beanie import Document, Indexed, Link


class DocumentA(Document):
    ...


class DocumentB(Document):
    a: Indexed(Link[DocumentA], unique=True)

Expected behavior

The code should run without issues, however, I get the following exception:

Show exception

Traceback (most recent call last):
File "/Users/edoardo/Library/Application Support/JetBrains/PyCharm2023.2/scratches/scratch.py", line 8, in <module>
  class DocumentB(Document):
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_model_construction.py", line 178, in __new__
  complete_model_class(
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_model_construction.py", line 468, in complete_model_class
  schema = cls.__get_pydantic_core_schema__(cls, handler)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/main.py", line 557, in __get_pydantic_core_schema__
  return __handler(__source)
         ^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_schema_generation_shared.py", line 82, in __call__
  schema = self._handler(__source_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 448, in generate_schema
  schema = self._generate_schema(obj)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 689, in _generate_schema
  return self._model_schema(obj)
         ^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 519, in _model_schema
  {k: self._generate_md_field_schema(k, v, decorators) for k, v in fields.items()},
  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 519, in <dictcomp>
  {k: self._generate_md_field_schema(k, v, decorators) for k, v in fields.items()},
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 853, in _generate_md_field_schema
  common_field = self._common_field_schema(name, field_info, decorators)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 905, in _common_field_schema
  schema = self._apply_annotations(
           ^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 1610, in _apply_annotations
  schema = get_inner_schema(source_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_schema_generation_shared.py", line 82, in __call__
  schema = self._handler(__source_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 1572, in inner_handler
  from_property = self._generate_schema_from_property(obj, obj)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 610, in _generate_schema_from_property
  schema = self._unpack_refs_defs(schema)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_generate_schema.py", line 567, in _unpack_refs_defs
  schema = flatten_schema_defs(schema)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 517, in flatten_schema_defs
  return _simplify_schema_references(schema, inline=False)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 432, in _simplify_schema_references
  schema = walk_core_schema(schema, collect_refs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 407, in walk_core_schema
  return f(schema, _dispatch)
         ^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 430, in collect_refs
  return recurse(s, collect_refs)
         ^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 183, in walk
  return f(schema.copy(), self._walk)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 430, in collect_refs
  return recurse(s, collect_refs)
         ^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 186, in _walk
  schema = self._schema_type_to_method[schema['type']](schema, f)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 195, in _handle_other_schemas
  schema['schema'] = self.walk(sub_schema, f)  # type: ignore
                     ^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 183, in walk
  return f(schema.copy(), self._walk)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 430, in collect_refs
  return recurse(s, collect_refs)
         ^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/edoardo/micromamba/envs/satml-llms-ctf/lib/python3.11/site-packages/pydantic/_internal/_core_utils.py", line 186, in _walk
  schema = self._schema_type_to_method[schema['type']](schema, f)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^
KeyError: 'Link'
/

All the exceptions seem to be happening in Pydantic, but it I guess this is happening because of how Indexed is created?

Additional context

I am on beanie==1.23.1 and pydantic==2.4.2, the latest version for both.

It looks like this should've been fixed by #754 but apparently it is not?

[dedeswim]

It seems to be related to (if not the same as) #701. Feel free to close if it is actually the same.

[roman-right]

Hi! Thank you for the issue, I'll check it during the next bug-fixing session next week.

[dedeswim]

Fixed by #762 if using Annotated. The following now runs:

from typing import Annotated

from beanie import Document, Indexed, Link


class DocumentA(Document):
    ...


class DocumentB(Document):
    a: Annotated[Link[DocumentA], Indexed(unique=True)]