Registry class¶

async def apply_default_force_nullable_fields(
    self,
    *,
    force_fields_nullable: Iterable[tuple[str, str]] | None = None,
    model_defaults: dict[str, dict[str, Any]] | None = None,
    filter_db_url: str | None = None,
    filter_db_name: str | None = None,
) -> None:
    """
    Applies default values to nullable fields in models, primarily used
    for online migrations.

    Args:
        force_fields_nullable (Iterable[tuple[str, str]] | None): A
            collection of (model_name, field_name) tuples for which default
            values should be applied. If None, uses values from
            `FORCE_FIELDS_NULLABLE` context variable.
        model_defaults (dict[str, dict[str, Any]] | None): A dictionary
            mapping model names to dictionaries of field_name: default_value
            pairs. These defaults will override existing field defaults.
        filter_db_url (str | None): If provided, only applies defaults to
                                    models connected to this specific
                                    database URL.
        filter_db_name (str | None): If provided, only applies defaults to
                                    models connected to this specific named
                                    extra database. Takes precedence over
                                    `filter_db_url`.
    """
    # Initialize force_fields_nullable from context var if not provided.
    if force_fields_nullable is None:
        force_fields_nullable = set(FORCE_FIELDS_NULLABLE.get())
    else:
        force_fields_nullable = set(force_fields_nullable)
    # Initialize model_defaults.
    if model_defaults is None:
        model_defaults = {}
    # Add model_defaults to force_fields_nullable.
    for model_name, defaults in model_defaults.items():
        for default_name in defaults:
            force_fields_nullable.add((model_name, default_name))
    # For empty model names, expand to include all matching models.
    for item in list(force_fields_nullable):
        if not item[0]:  # If model name is empty string.
            force_fields_nullable.discard(item)
            for model in self.models.values():
                if item[1] in model.meta.fields:
                    force_fields_nullable.add((model.__name__, item[1]))

    if not force_fields_nullable:
        return  # No fields to process, exit early.

    # Determine the database URL to filter by.
    if isinstance(filter_db_name, str):
        if filter_db_name:
            filter_db_url = str(self.extra[filter_db_name].url)
        else:
            filter_db_url = str(self.database.url)

    models_with_fields: dict[str, set[str]] = {}
    # Populate models_with_fields with models and their relevant fields.
    for item in force_fields_nullable:
        if item[0] not in self.models:
            continue
        if item[1] not in self.models[item[0]].meta.fields:
            continue
        # Check if field has a default or if an override default is provided.
        if not self.models[item[0]].meta.fields[item[1]].has_default():
            overwrite_default = model_defaults.get(item[0]) or {}
            if item[1] not in overwrite_default:
                continue
        field_set = models_with_fields.setdefault(item[0], set())
        field_set.add(item[1])

    if not models_with_fields:
        return  # No valid models/fields to update, exit early.

    ops = []
    # Iterate through models and their fields to create update operations.
    for model_name, field_set in models_with_fields.items():
        model = self.models[model_name]
        if filter_db_url and str(model.database.url) != filter_db_url:
            continue  # Skip if database URL does not match filter.

        model_specific_defaults = model_defaults.get(model_name) or {}
        filter_kwargs = dict.fromkeys(field_set)

        async def wrapper_fn(
            _model: type[BaseModelType] = model,
            _model_specific_defaults: dict = model_specific_defaults,
            _filter_kwargs: dict = filter_kwargs,
            _field_set: set[str] = field_set,
        ) -> None:
            # To reduce memory usage, only retrieve pknames and load per object.
            # Load all at once to prevent cursor interference with updates.
            query = _model.query.filter(**_filter_kwargs).only(*_model.pknames)
            for obj in await query:
                await obj.load()  # Load the full object data.
                # Extract database fields, excluding those in _field_set.
                kwargs = {
                    k: v for k, v in obj.extract_db_fields().items() if k not in _field_set
                }
                kwargs.update(_model_specific_defaults)  # Apply specific defaults.
                # We serialize per table to avoid transaction interlocking errors.
                # Also, tables can become very large.
                token = CURRENT_INSTANCE.set(query)
                try:
                    await obj._update(
                        False,  # is_partial=False
                        kwargs,
                        # Pre-update signal for migrations.
                        pre_fn=partial(
                            _model.meta.signals.pre_update.send_async,
                            is_update=True,
                            is_migration=True,
                        ),
                        # Post-update signal for migrations.
                        post_fn=partial(
                            _model.meta.signals.post_update.send_async,
                            is_update=True,
                            is_migration=True,
                        ),
                        instance=query,
                    )
                finally:
                    CURRENT_INSTANCE.reset(token)  # Reset context variable.

        ops.append(wrapper_fn())

    # Execute all update operations concurrently.
    await run_concurrently(ops, limit=settings.orm_registry_ops_limit)

def extra_name_check(self, name: Any) -> bool:
    """
    Validates the name of an extra database connection.

    Args:
        name (Any): The name to validate.

    Returns:
        bool: True if the name is valid, False otherwise. Logs errors/warnings.
    """
    if not isinstance(name, str):
        logger.error(f"Extra database name: {name!r} is not a string.")
        return False
    elif not name.strip():
        logger.error(f'Extra database name: "{name}" is empty.')
        return False

    if name.strip() != name:
        logger.warning(
            f'Extra database name: "{name}" starts or ends with whitespace characters.'
        )
    return True

def __copy__(self) -> Registry:
    """
    Creates a shallow copy of the Registry instance, including its models
    and their metadata.

    Returns:
        Registry: A new Registry instance with copied models.
    """
    content_type: bool | type[BaseModelType] = False
    if self.content_type is not None:
        try:
            # Attempt to copy the ContentType model if it exists and is copyable.
            content_type = self.get_model(
                "ContentType", include_content_type_attr=False
            ).copy_edgy_model()
        except LookupError:
            # Fallback to the original content_type if not found.
            content_type = self.content_type
    # Create a new Registry instance with basic settings.
    _copy = Registry(
        self.database, with_content_type=content_type, schema=self.db_schema, extra=self.extra
    )
    # Copy models from different registry types.
    for registry_type in self.model_registry_types:
        dict_models = getattr(_copy, registry_type)
        dict_models.update(
            (
                (
                    key,
                    val.copy_edgy_model(registry=_copy),
                )
                for key, val in getattr(self, registry_type).items()
                if not val.meta.no_copy and key not in dict_models
            )
        )
    _copy.dbs_reflected = set(self.dbs_reflected)  # Copy reflected databases.
    return _copy

def _set_content_type(
    self,
    with_content_type: Literal[True] | type[BaseModelType],
    old_content_type_to_replace: type[BaseModelType] | None = None,
) -> None:
    """
    Configures content type support within the registry. This involves
    either creating a default ContentType model or registering a provided
    one, and then setting up callbacks to automatically add a 'content_type'
    field to other models.

    Args:
        with_content_type (Literal[True] | type[BaseModelType]): If True,
            uses the default Edgy ContentType model. If a BaseModelType,
            that model will be used as the ContentType.
        old_content_type_to_replace (type[BaseModelType] | None): An
            optional existing ContentType model that needs to be replaced
            with the new one (e.g., during registry copying).
    """
    from edgy.contrib.contenttypes.fields import BaseContentTypeField, ContentTypeField
    from edgy.contrib.contenttypes.models import ContentType
    from edgy.core.db.models.metaclasses import MetaInfo
    from edgy.core.db.relationships.related_field import RelatedField
    from edgy.core.utils.models import create_edgy_model

    # Use default ContentType if `with_content_type` is True.
    if with_content_type is True:
        with_content_type = ContentType

    real_content_type: type[BaseModelType] = with_content_type

    # If the provided content type model is abstract, create a concrete one.
    if real_content_type.meta.abstract:
        in_admin = real_content_type.meta.in_admin
        no_admin_create = real_content_type.meta.no_admin_create
        meta_args = {
            "tablename": "contenttypes",
            "registry": self,
            "in_admin": True if in_admin is None else in_admin,
            "no_admin_create": True if no_admin_create is None else no_admin_create,
        }

        new_meta: MetaInfo = MetaInfo(None, **meta_args)
        # Model adds itself to registry and executes callbacks.
        real_content_type = create_edgy_model(
            "ContentType",
            with_content_type.__module__,
            __metadata__=new_meta,
            __bases__=(with_content_type,),
        )
    # If the content type model is not abstract but not yet in this registry.
    elif real_content_type.meta.registry is None:
        real_content_type.add_to_registry(self, name="ContentType")
    self.content_type = real_content_type

    def callback(model_class: type[BaseModelType]) -> None:
        """
        Callback function executed when a model is registered. It adds a
        'content_type' field to the model if not already present.
        """
        # Skip ContentType model itself to avoid recursion.
        if issubclass(model_class, ContentType):
            return
        # Skip if field is explicitly set or removed when copying.
        for field in model_class.meta.fields.values():
            if isinstance(field, BaseContentTypeField):
                if (
                    old_content_type_to_replace is not None
                    and field.target is old_content_type_to_replace
                ):
                    field.target_registry = self
                    field.target = real_content_type
                    # Simply overwrite the related field in ContentType.
                    real_content_type.meta.fields[field.related_name] = RelatedField(
                        name=field.related_name,
                        foreign_key_name=field.name,
                        related_from=model_class,
                        owner=real_content_type,
                    )
                return

        # E.g., if the field is explicitly excluded.
        if "content_type" in model_class.meta.fields:
            return

        related_name = f"reverse_{model_class.__name__.lower()}"
        # Ensure no duplicate related name in ContentType.
        assert related_name not in real_content_type.meta.fields, (
            f"duplicate model name: {model_class.__name__}"
        )

        field_args: dict[str, Any] = {
            "name": "content_type",
            "owner": model_class,
            "to": real_content_type,
            "no_constraint": real_content_type.no_constraint,
            "no_copy": True,
        }
        # Set cascade deletion properties if registries differ.
        if model_class.meta.registry is not real_content_type.meta.registry:
            field_args["relation_has_post_delete_callback"] = True
            field_args["force_cascade_deletion_relation"] = True
        # Add the ContentTypeField to the model.
        model_class.meta.fields["content_type"] = ContentTypeField(**field_args)
        # Add the reverse related field to ContentType.
        real_content_type.meta.fields[related_name] = RelatedField(
            name=related_name,
            foreign_key_name="content_type",
            related_from=model_class,
            owner=real_content_type,
        )

    # Register the callback to be executed for all models (not one-time).
    self.register_callback(None, callback, one_time=False)

def get_model(
    self,
    model_name: str,
    *,
    include_content_type_attr: bool = True,
    exclude: Container[str] = (),
) -> type[BaseModelType]:
    """
    Retrieves a registered model by its name.

    Args:
        model_name (str): The name of the model to retrieve.
        include_content_type_attr (bool): If True and `model_name` is
            "ContentType", returns the configured content type model.
            Defaults to True.
        exclude (Container[str]): A collection of registry types (e.g.,
                                  "pattern_models") to exclude from the
                                  search.

    Returns:
        type[BaseModelType]: The found model class.

    Raises:
        LookupError: If no model with the given name is found.
    """
    # Handle special case for ContentType model.
    if (
        include_content_type_attr
        and model_name == "ContentType"
        and self.content_type is not None
    ):
        return self.content_type
    # Search through various model registries.
    for model_dict_name in self.model_registry_types:
        if model_dict_name in exclude:
            continue
        model_dict: dict = getattr(self, model_dict_name)
        if model_name in model_dict:
            return cast(type["BaseModelType"], model_dict[model_name])
    raise LookupError(f'Registry doesn\'t have a "{model_name}" model.') from None

def delete_model(self, model_name: str) -> bool:
    """
    Deletes a model from the registry by its name.

    Args:
        model_name (str): The name of the model to delete.

    Returns:
        bool: True if the model was found and deleted, False otherwise.
    """
    self.admin_models.discard(model_name)  # Remove from admin models set.
    for model_dict_name in self.model_registry_types:
        model_dict: dict = getattr(self, model_dict_name)
        if model_name in model_dict:
            del model_dict[model_name]  # Delete the model.
            return True
    return False

def refresh_metadata(
    self,
    *,
    update_only: bool = False,
    multi_schema: bool | re.Pattern | str = False,
    ignore_schema_pattern: re.Pattern | str | None = "information_schema",
    _schemes_tree: dict[str | None, tuple[str, sqlalchemy.MetaData, list[str]]] | None = None,
) -> None:
    """
    Refreshes the SQLAlchemy MetaData objects associated with the models
    in the registry. This is crucial for ensuring that table definitions
    are up-to-date, especially in multi-schema or dynamic reflection scenarios.

    Args:
        update_only (bool): If True, only updates existing table definitions
                            without clearing the metadata first. Defaults to
                            False.
        multi_schema (bool | re.Pattern | str): If True, enables multi-schema
            reflection based on detected schemas. Can also be a regex pattern
            or string to match specific schemas. Defaults to False.
        ignore_schema_pattern (re.Pattern | str | None): A regex pattern
            or string to ignore certain schemas during multi-schema reflection.
            Defaults to "information_schema".
    """
    if not update_only:
        for val in self.metadata_by_name.values():
            val.clear()  # Clear existing metadata if not just updating.

    maindatabase_url = str(self.database.url)
    # Determine schemes to process based on multi_schema setting.
    if multi_schema is not False:
        if _schemes_tree is None:
            _schemes_tree = run_sync(self.schema.get_schemes_tree(no_reflect=True))
        schemes_tree: dict[str, tuple[str | None, list[str]]] = {
            v[0]: (key, v[2]) for key, v in _schemes_tree.items()
        }
    else:
        schemes_tree = {
            # main db with db_schema
            maindatabase_url: (None, [self.db_schema]),
            # extra dbs without
            **{str(v.url): (k, [None]) for k, v in self.extra.items()},
        }

    # Compile regex patterns if provided as strings.
    if isinstance(multi_schema, str):
        multi_schema = re.compile(multi_schema)
    if isinstance(ignore_schema_pattern, str):
        ignore_schema_pattern = re.compile(ignore_schema_pattern)

    # clear data
    if not update_only:
        for model_class in chain(
            self.tenant_models.values(), self.models.values(), self.reflected.values()
        ):
            model_class._table = None  # Clear cached table.
            model_class._db_schemas = {}  # Clear cached db schemas.
    db_schema = self.db_schema or ""
    # Iterate through all registered models.
    # In case of models which double as tenant models, we check the regular case
    for model_class in self.models.values():
        url = str(model_class.database.url)
        if url in schemes_tree:
            extra_key, schemes = schemes_tree[url]
            for schema in schemes:
                if multi_schema is not False:
                    assert schema is not None
                    # Skip if multi_schema is enabled but pattern doesn't match.
                    if multi_schema is not True and multi_schema.match(schema) is None:
                        continue
                    # Skip if schema matches ignore pattern.
                    if (
                        ignore_schema_pattern is not None
                        and ignore_schema_pattern.match(schema) is not None
                    ):
                        continue
                    # Handle model schemes
                    # Case 1: no schema defined
                    if (
                        model_class.__using_schema__ is Undefined
                        or model_class.__using_schema__ is None
                    ):
                        # if not main schema, continue
                        if schema != db_schema:
                            continue
                    elif model_class.__using_schema__ != schema:
                        # Case 2: non-matching schema, continue
                        continue

                # Initialize table schema for the model.
                model_class.table_schema(
                    schema=schema, metadata=self.metadata_by_url[url], update_cache=True
                )
                break

    # Iterate through all registered tenant models (not necessarily in models).
    for model_class in self.tenant_models.values():
        url = str(model_class.database.url)
        if url in schemes_tree:
            extra_key, schemes = schemes_tree[url]
            for schema in schemes:
                if multi_schema is not False:
                    assert schema is not None
                    # Skip if multi_schema is enabled but pattern doesn't match.
                    if multi_schema is not True and multi_schema.match(schema) is None:
                        continue
                    # Skip if schema matches ignore pattern.
                    if (
                        ignore_schema_pattern is not None
                        and ignore_schema_pattern.match(schema) is not None
                    ):
                        continue
                # Initialize table schema for the model.
                model_class.table_schema(
                    schema=schema, metadata=self.metadata_by_url[url], update_cache=True
                )

async def arefresh_metadata(
    self,
    *,
    update_only: bool = False,
    multi_schema: bool | re.Pattern | str = False,
    ignore_schema_pattern: re.Pattern | str | None = "information_schema",
) -> None:
    """
    Refreshes the SQLAlchemy MetaData objects associated with the models
    in the registry. This is crucial for ensuring that table definitions
    are up-to-date, especially in multi-schema or dynamic reflection scenarios.

    Args:
        update_only (bool): If True, only updates existing table definitions
                            without clearing the metadata first. Defaults to
                            False.
        multi_schema (bool | re.Pattern | str): If True, enables multi-schema
            reflection based on detected schemas. Can also be a regex pattern
            or string to match specific schemas. Defaults to False.
        ignore_schema_pattern (re.Pattern | str | None): A regex pattern
            or string to ignore certain schemas during multi-schema reflection.
            Defaults to "information_schema".

    Difference to refresh_metadata: if multi_schema is in use, don't use a second loop
    """
    _schemes_tree = (
        None if multi_schema is False else await self.schema.get_schemes_tree(no_reflect=True)
    )
    self.refresh_metadata(
        update_only=update_only,
        multi_schema=multi_schema,
        ignore_schema_pattern=ignore_schema_pattern,
        _schemes_tree=_schemes_tree,
    )

def register_callback(
    self,
    name_or_class: type[BaseModelType] | str | None,
    callback: Callable[[type[BaseModelType]], None],
    one_time: bool | None = None,
) -> None:
    """
    Registers a callback function to be executed when a model is added
    or a specific model is accessed.

    Args:
        name_or_class (type[BaseModelType] | str | None): The model class,
            model name (string), or None for a general callback applied to
            all models.
        callback (Callable[[type[BaseModelType]], None]): The callback
                                                          function to
                                                          execute. It takes
                                                          the model class
                                                          as an argument.
        one_time (bool | None): If True, the callback will only be executed
                                once. If None, it defaults to True for
                                model-specific callbacks and False for
                                general callbacks.
    """
    if one_time is None:
        # True for model specific callbacks, False for general callbacks.
        one_time = name_or_class is not None
    called: bool = False
    if name_or_class is None:  # General callback for all models.
        for model in self.models.values():
            callback(model)
            called = True
        for model in self.reflected.values():
            callback(model)
            called = True
        for name, model in self.tenant_models.items():
            # For tenant-only models, ensure they are not already in general models.
            if name not in self.models:
                callback(model)
                called = True
    elif not isinstance(name_or_class, str):  # Specific model class.
        callback(name_or_class)
        called = True
    else:  # Specific model by name.
        model_class = None
        with contextlib.suppress(LookupError):
            model_class = self.get_model(name_or_class)
        if model_class is not None:
            callback(model_class)
            called = True
    # Convert model class to its name if it was passed as a type.
    if name_or_class is not None and not isinstance(name_or_class, str):
        name_or_class = name_or_class.__name__
    if called and one_time:
        return  # If already called and is one-time, exit.
    if one_time:
        self._onetime_callbacks[name_or_class].append(callback)
    else:
        self._callbacks[name_or_class].append(callback)

def execute_model_callbacks(self, model_class: type[BaseModelType]) -> None:
    """
    Executes all registered callbacks (one-time and persistent) for a
    given model class.

    Args:
        model_class (type[BaseModelType]): The model class for which to
                                           execute callbacks.
    """
    name = model_class.__name__
    # Execute one-time callbacks specific to this model.
    callbacks = self._onetime_callbacks.get(name)
    while callbacks:
        callbacks.pop()(model_class)

    # Execute general one-time callbacks.
    callbacks = self._onetime_callbacks.get(None)
    while callbacks:
        callbacks.pop()(model_class)

    # Execute persistent callbacks specific to this model.
    callbacks = self._callbacks.get(name)
    if callbacks:
        for callback in callbacks:
            callback(model_class)

    # Execute general persistent callbacks.
    callbacks = self._callbacks.get(None)
    if callbacks:
        for callback in callbacks:
            callback(model_class)

def init_models(
    self, *, init_column_mappers: bool = True, init_class_attrs: bool = True
) -> None:
    """
    Initializes lazy-loaded parts of model metadata (e.g., column mappers
    and class attributes). This method is normally not required to be called
    explicitly as it's handled internally.

    Args:
        init_column_mappers (bool): If True, initializes SQLAlchemy column
                                    mappers. Defaults to True.
        init_class_attrs (bool): If True, initializes model class attributes.
                                 Defaults to True.
    """
    for model_class in self.models.values():
        model_class.meta.full_init(
            init_column_mappers=init_column_mappers, init_class_attrs=init_class_attrs
        )

    for model_class in self.reflected.values():
        model_class.meta.full_init(
            init_column_mappers=init_column_mappers, init_class_attrs=init_class_attrs
        )

def invalidate_models(self, *, clear_class_attrs: bool = True) -> None:
    """
    Invalidates all lazy-loaded parts of model metadata. They will be
    automatically re-initialized upon next access. This is useful for
    scenarios where model definitions might change dynamically.

    Args:
        clear_class_attrs (bool): If True, clears cached class attributes.
                                  Defaults to True.
    """
    for model_class in self.models.values():
        model_class.meta.invalidate(clear_class_attrs=clear_class_attrs)
    for model_class in self.reflected.values():
        model_class.meta.invalidate(clear_class_attrs=clear_class_attrs)

def get_tablenames(self) -> set[str]:
    """
    Returns a set of all table names associated with the models registered
    in this registry (including reflected models).
    """
    return_set = set()
    for model_class in self.models.values():
        return_set.add(model_class.meta.tablename)
    for model_class in self.reflected.values():
        return_set.add(model_class.meta.tablename)
    return return_set

def _automigrate_update(
    self,
    migration_settings: EdgySettings,
) -> None:
    """
    Internal synchronous method to run database migrations using Monkay.

    Args:
        migration_settings (EdgySettings): Settings specific to the migration
                                          process.
    """
    from edgy import Instance, monkay
    from edgy.cli.base import upgrade

    self._is_automigrated = True
    with monkay.with_full_overwrite(
        extensions={},
        settings=migration_settings,
        instance=Instance(registry=self),
        evaluate_settings_with={},
        apply_extensions=True,
    ):
        upgrade()