def model_config(shard_group=None, database=None, skip_runtime_checks=False):
    """
    A decorator for marking a model as being either sharded or stored on a
    particular database. When sharding, it does some verification to ensure
    that the model is defined correctly.
    """
    def configure(cls):
        if database and shard_group:
            raise ShardedModelInitializationException('A model cannot be both sharded and stored on a particular database.')

        if not database and not shard_group:
            raise ShardedModelInitializationException('The model should be either sharded or stored on a database in the `model_config` decorator is used.')

        if database:
            if (database not in settings.DATABASES or settings.DATABASES[database].get('PRIMARY')) and skip_runtime_checks is False:
                raise NonExistentDatabaseException(
                    'Unable to place {} in {} as that is not an existing primary database in the system.'.format(cls._meta.model_name, database)
                )
            setattr(cls, 'django_sharding__database', database)

        postgres_shard_id_fields = list(filter(lambda field: issubclass(type(field), BasePostgresShardGeneratedIDField), cls._meta.fields))
        if postgres_shard_id_fields and skip_runtime_checks is False:
            database_dicts = [settings.DATABASES[database]] if database else [db_settings for db, db_settings in
                                                                              iteritems(settings.DATABASES) if
                                                                              db_settings["SHARD_GROUP"] == shard_group]
            if any([database_dict['ENGINE'] not in Backends.POSTGRES for database_dict in database_dicts]):
                raise ShardedModelInitializationException(
                    'You cannot use a PostgresShardGeneratedIDField on a non-Postgres database.')

            for field in postgres_shard_id_fields:
                register_migration_signal_for_model_receiver(apps.get_app_config(cls._meta.app_label),
                                                             field.migration_receiver,
                                                             dispatch_uid=PRE_MIGRATION_DISPATCH_UID % cls._meta.app_label)

        if shard_group:
            sharded_fields = list(filter(lambda field: issubclass(type(field), ShardedIDFieldMixin), cls._meta.fields))
            if not sharded_fields and not postgres_shard_id_fields:
                raise ShardedModelInitializationException('All sharded models require a ShardedIDFieldMixin or a '
                                                          'PostgresShardGeneratedIDField.')

            if not list(filter(lambda field: field == cls._meta.pk, sharded_fields + postgres_shard_id_fields)):
                raise ShardedModelInitializationException('All sharded models require a ShardedAutoIDField or '
                                                          'PostgresShardGeneratedIDField to be the primary key. Set '
                                                          'primary_key=True on the field.')

            if not callable(getattr(cls, 'get_shard', None)):
                raise ShardedModelInitializationException('You must define a get_shard method on the sharded model.')

            setattr(cls, 'django_sharding__shard_group', shard_group)
            setattr(cls, 'django_sharding__is_sharded', True)

        return cls
    return configure def shard_storage_config(shard_group='default', shared_field='shard'):
    def configure(cls):
        setattr(cls, 'django_sharding__shard_group', shard_group)
        setattr(cls, 'django_sharding__shard_field', shared_field)
        setattr(cls, 'django_sharding__stores_shard', True)
        return cls
    return configure class TableShardedIDField(ShardedIDFieldMixin, BigAutoField):
    """
    An autoincrimenting field which takes a `source_table_name` as an argument in
    order to generate unqiue ids for the sharded model.  i.e. `app.model_name`.
    """
    def __init__(self, *args, **kwargs):
        from django_sharding_library.id_generation_strategies import TableStrategy
        kwargs['strategy'] = TableStrategy(backing_model_name=kwargs['source_table_name'])
        setattr(self, 'source_table_name', kwargs['source_table_name'])
        del kwargs['source_table_name']
        return super(TableShardedIDField, self).__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super(TableShardedIDField, self).deconstruct()
        kwargs['source_table_name'] = getattr(self, 'source_table_name')
        return name, path, args, kwargs class ShardForeignKeyStorageField(ShardForeignKeyStorageFieldMixin, ForeignKey):
    """
    A simple char field that stores a shard and uses a signal to generate
    the shard using a pre_save signal.
    """
    pass class PostgresShardGeneratedIDAutoField(BasePostgresShardGeneratedIDField, BigAutoField):
    """
    A field that uses a Postgres stored procedure to return an ID generated on the database.
    """
    def db_type(self, connection, *args, **kwargs):
        try:
            from django.db.backends.postgresql.base import DatabaseWrapper as PostgresDatabaseWrapper
        except ImportError:
            from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as PostgresDatabaseWrapper

        if connection.vendor == PostgresDatabaseWrapper.vendor:
            return "bigint DEFAULT next_sharded_id()"
        else:
            return super(PostgresShardGeneratedIDAutoField, self).db_type(connection) class PostgresShardGeneratedIDField(BasePostgresShardGeneratedIDField, BigIntegerField):
    """
    A field that uses a Postgres stored procedure to return an ID generated on the database.

    Generates them prior to save with a seperate call to the DB.
    """

    def get_shard_from_id(self, instance_id):
        group = getattr(self, 'django_sharding__shard_group', None)
        shard_id_to_find = int(bin(instance_id)[-23:-10], 2)  # We know where the shard id is stored in the PK's bits.

        # We can check the shard id from the PK against the shard ID in the databases config
        for alias, db_settings in settings.DATABASES.items():
            if db_settings["SHARD_GROUP"] == group and db_settings["SHARD_ID"] == shard_id_to_find:
                return alias

        return None  # Return None if we could not determine the shard so we can fall through to the next shard grab attempt

    def get_pk_value_on_save(self, instance):
        return self.generate_id(instance)

    def pre_save(self, model_instance, add):
        if getattr(model_instance, self.attname, None) is not None:
            return super(PostgresShardGeneratedIDField, self).pre_save(model_instance, add)
        value = self.generate_id(model_instance)
        setattr(model_instance, self.attname, value)
        return value

    @staticmethod
    def generate_id(instance):
        shard = instance._state.db or instance.get_shard()
        return get_next_sharded_id(shard) class ShardedByMixin(models.Model):
    """
    This mixin is intended to be used with the included `save_shard_handler` in order
    to save the shard to the object. This is done by registering it on the User's pre-save signal.
    Look at the signal's docstring for additional information.

    Note: requires the use of `cls.shard_group`
    """
    django_sharding__shard_field = 'shard'
    django_sharding__stores_shard = True

    SHARD_CHOICES = ((i, i) for i in _get_primary_shards())

    shard = models.CharField(max_length=120, blank=True, null=True, choices=SHARD_CHOICES)

    class Meta:
        abstract = True class ShardStorageModel(models.Model):
    """
    A model for storing shards.
    """
    SHARD_CHOICES = ((i, i) for i in _get_primary_shards())

    shard = models.CharField(max_length=120, choices=SHARD_CHOICES)
    shard_key = models.CharField(primary_key=True, max_length=120)

    class Meta:
        abstract = True class TableStrategyModel(models.Model):
    """
    A model for use with the TableStrategy in order to generate ids using an
    autoincrimenting field. To use this, simply inherit from this class and give
    the model a more meaningful name e.g. InvoiceIDs.
    """
    id = BigAutoField(primary_key=True)
    stub = models.NullBooleanField(null=True, default=True, unique=True)

    class Meta:
        abstract = True class Backends(object):
    MYSQL = ('django.db.backends.mysql', 'django.contrib.gis.db.backends.mysql')
    POSTGRES = ('django.db.backends.postgresql_psycopg2', 'django.db.backends.postgresql', 'django.contrib.gis.db.backends.postgis')
    SQLITE = ('django.db.backends.sqlite3', 'django.contrib.gis.db.backends.spacialite')
from django.contrib.auth.models import AbstractUser
from django.db import models
from django.conf import settings
from django_sharding_library.decorators import model_config, shard_storage_config
from django_sharding_library.fields import (
    TableShardedIDField,
    ShardForeignKeyStorageField,
    PostgresShardGeneratedIDAutoField,
    PostgresShardGeneratedIDField
)
from django_sharding_library.models import ShardedByMixin, ShardStorageModel, TableStrategyModel
from django_sharding_library.constants import Backends
        from django.contrib.auth import get_user_model
        from django.contrib.auth import get_user_model
        from django.contrib.auth import get_user_model


# A model for use with a sharded model to generate pk's using
# an autoincrement field on the backing TableStrategyModel.
# This one is initialized but used in no strategy.
class ShardedModelIDs(TableStrategyModel):
    pass


# An implementation of the extension of a the Django user to add
# the mixin provided in order to save the shard on the user.
@shard_storage_config()
class User(AbstractUser, ShardedByMixin):
    # Patch for Django 1.11
    groups = None
    user_permissions = None


# An implementation of the extension of a the Django user to add
# the mixin provided in order to save the shard on the user.
@shard_storage_config(shard_group='postgres')

class PostgresShardUser(AbstractUser, ShardedByMixin):