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') class ShardedIDFieldMixin(object):
    """
    An autoincrimenting field which takes an id generator class instance
    as an argument and uses the generator to assign each new object a unique
    id.
    Note: This currently must be the primary key of the model and although this
    may be updated in the future, it should not hinder the app to use other
    candidates as unique fields.
    """
    def __init__(self, *args, **kwargs):
        # Remove the strategy from the kwargs so that it doesn't get passed to Django.
        setattr(self, 'strategy', kwargs['strategy'])
        del kwargs['strategy']
        return super(ShardedIDFieldMixin, self).__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super(ShardedIDFieldMixin, self).deconstruct()

        # Add the strategy from the kwargs so that it does get passed to our model.
        kwargs['strategy'] = getattr(self, 'strategy')
        return name, path, args, kwargs

    def get_pk_value_on_save(self, instance):
        if not instance.pk:
            return self.strategy.get_next_id()
        return instance.pk class ShardLocalStorageFieldMixin(ShardStorageFieldMixin):
    """
    The ShardLocalStorageFieldMixin is used for when the shard is stored on the model that
    is being sharded by. i.e. Storing the shard on the User model and sharding by the User.
    """
    def __init__(self, *args, **kwargs):
        setattr(self, 'django_sharding__use_signal', True)
        return super(ShardLocalStorageFieldMixin, self).__init__(*args, **kwargs)

    def deconstruct(self):
        return super(ShardLocalStorageFieldMixin, self).deconstruct() class ShardStorageFieldMixin(object):
    """
    A mixin for a field used to store a shard for in an instance or parent of an instance.
    """
    def __init__(self, *args, **kwargs):
        setattr(self, 'django_sharding__stores_shard', True)
        setattr(self, 'django_sharding__shard_group', kwargs['shard_group'])
        del kwargs['shard_group']
        return super(ShardStorageFieldMixin, self).__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super(ShardStorageFieldMixin, self).deconstruct()
        kwargs['shard_group'] = getattr(self, 'django_sharding__shard_group')
        return name, path, args, kwargs class ShardForeignKeyStorageFieldMixin(ShardStorageFieldMixin):
    """
    A mixin for a field used to store a foreign key to another table which
    stores the shard, usually a table which inherits from the ShardStorageModel.
    """
    def __init__(self, *args, **kwargs):
        setattr(self, 'django_sharding__stores_shard', True)
        model_class = kwargs.get('to', args and args[0])
        if type(model_class) == str:
            app_label = model_class.split('.')[0]
            app = apps.get_app_config(app_label)
            model_class = app.get_model(model_class[len(app_label) + 1:], require_ready=False)
        setattr(self, 'django_sharding__shard_storage_table', model_class)
        return super(ShardForeignKeyStorageFieldMixin, self).__init__(*args, **kwargs)

    def deconstruct(self):
        name, path, args, kwargs = super(ShardForeignKeyStorageFieldMixin, self).deconstruct()
        return name, path, args, kwargs

    def pre_save(self, model_instance, add):
        self.save_shard(model_instance)
        return super(ShardForeignKeyStorageFieldMixin, self).pre_save(model_instance, add)

    def save_shard(self, model_instance):
        shard_key = model_instance.get_shard_key()
        if not getattr(model_instance, self.name):
            shard_storage_table = getattr(self, 'django_sharding__shard_storage_table')
            shard_group = getattr(self, 'django_sharding__shard_group')

            app_config_app_label = getattr(settings, 'DJANGO_SHARDING_SETTINGS', {}).get('APP_CONFIG_APP', 'django_sharding')
            bucketer = apps.get_app_config(app_config_app_label).get_bucketer(shard_group)
            shard = bucketer.pick_shard(model_instance)
            shard_object, _ = shard_storage_table.objects.get_or_create(shard_key=shard_key)
            if not shard_object.shard:
                shard_object.shard = shard
                shard_object.save()
            setattr(model_instance, self.name, shard_object) 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 BaseIDGenerationStrategy(object):
    """
    A strategy for Generating unique identifiers for the sharded models.
    """
    def get_next_id(self, database=None):
        """
        A function which takes no arguments and returns a new unique identifier.
        """
        raise NotImplementedError class ShardedModelIDs(TableStrategyModel):
    pass class ShardedTestModelIDs(TableStrategyModel):
    pass class TestModel(models.Model):
    id = TableShardedIDField(primary_key=True, source_table_name='tests.ShardedTestModelIDs')
    random_string = models.CharField(max_length=120)
    user_pk = models.PositiveIntegerField()

    def get_shard(self):
        from django.contrib.auth import get_user_model
        return get_user_model().objects.get(pk=self.user_pk).shard

    @staticmethod
    def get_shard_from_id(user_pk):
        from django.contrib.auth import get_user_model
        user = get_user_model()
        return user.objects.get(pk=user_pk).shard class ShardStorageTable(ShardStorageModel):
    pass class PostgresCustomAutoIDModel(models.Model):
    if settings.DATABASES['default']['ENGINE'] in Backends.POSTGRES:
        id = PostgresShardGeneratedIDAutoField(primary_key=True)
    else:
        id = TableShardedIDField(primary_key=True, source_table_name='tests.PostgresCustomIDModelBackupField')
    random_string = models.CharField(max_length=120)
    user_pk = models.PositiveIntegerField()

    def get_shard(self):
        return PostgresShardUser.objects.get(pk=self.user_pk).shard

    @staticmethod
    def get_shard_from_id(user_pk):
        return PostgresShardUser.objects.get(pk=user_pk).shard class PostgresCustomIDModel(models.Model):
    if settings.DATABASES['default']['ENGINE'] in Backends.POSTGRES:
        id = PostgresShardGeneratedIDField(primary_key=True)
        some_field = PostgresShardGeneratedIDField()
    else:
        id = TableShardedIDField(primary_key=True, source_table_name='tests.PostgresCustomIDModelBackupField')
        some_field = models.PositiveIntegerField()
    random_string = models.CharField(max_length=120)
    user_pk = models.PositiveIntegerField()

    def get_shard(self):
        return PostgresShardUser.objects.get(pk=self.user_pk).shard

    @staticmethod
    def get_shard_from_id(user_pk):
        return PostgresShardUser.objects.get(pk=user_pk).shard class PostgresShardUser(AbstractUser, ShardedByMixin):
    # Patch for Django 1.11
    groups = None
    user_permissions = None
from mock import patch
from six.moves import xrange
from django.conf import settings
from django.db import connections
from django.db.utils import DataError, IntegrityError
from django.db.models import CASCADE
from django.test import TestCase
from datetime import datetime
from django_sharding_library.constants import Backends
from django_sharding_library.fields import (
    ShardedIDFieldMixin,
    ShardLocalStorageFieldMixin,
    ShardStorageFieldMixin,
    ShardForeignKeyStorageFieldMixin,
    ShardForeignKeyStorageField,
)
from django_sharding_library.id_generation_strategies import BaseIDGenerationStrategy
from tests.models import (
    ShardedModelIDs,
    ShardedTestModelIDs,
    TestModel,
    ShardStorageTable,
    PostgresCustomAutoIDModel,
    PostgresCustomIDModel,
    PostgresShardUser
)
import unittest
import time





class BigAutoFieldTestCase(TestCase):
    databases = '__all__'

    def test_largest_id(self):
        if settings.DATABASES['default']['ENGINE'] in Backends.POSTGRES + Backends.SQLITE:
            max_id = 9223372036854775807
        else:
            max_id = 18446744073709551615

        item = ShardedModelIDs.objects.create(id=max_id, stub=None)
        self.assertEqual(item.pk, max_id)
        try:
            ShardedModelIDs.objects.create(id=max_id + 1, stub=None)
        except (DataError, OverflowError, IntegrityError):
            pass
        else:
            raise Exception('Field was able to hold an integer as large as {}.'.format(max_id + 1))


class FakeIDGenerationStrategy(BaseIDGenerationStrategy):
    def __init__(self):
        self.call_count = 0
        self.id = 0

    def get_next_id(self):
        self.call_count += 1
        self.id += 1
        return self.id


class TableShardedIDFieldTestCase(TestCase):
    databases = '__all__'

    def test_largest_id(self):
        if settings.DATABASES['default']['ENGINE'] in Backends.POSTGRES + Backends.SQLITE:
            max_id = 9223372036854775807
        else:
            max_id = 18446744073709551615


        item = TestModel.objects.using('app_shard_001').create(id=max_id, user_pk=1)