# Path: django_sharding_library/constants.py
# 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')
# 
# Path: django_sharding_library/utils.py
# def create_postgres_global_sequence(sequence_name, db_alias, reset_sequence=False):
#     cursor = connections[db_alias].cursor()
#     sid = transaction.savepoint(db_alias)
#     create_sequence_if_not_exists_sql = """DO
# $$
# BEGIN
#         CREATE SEQUENCE %s;
# EXCEPTION WHEN duplicate_table THEN
#         -- do nothing, it's already there
# END
# $$ LANGUAGE plpgsql;"""
#     try:
#         cursor.execute(create_sequence_if_not_exists_sql % sequence_name)
#     except DatabaseError:
#         transaction.savepoint_rollback(sid, using=db_alias)
#     else:
#         transaction.savepoint_commit(sid, using=db_alias)
#     if reset_sequence:
#         cursor.execute("SELECT setval('%s', 1, false)" % (sequence_name,))
#     cursor.close()
# 
# def create_postgres_shard_id_function(sequence_name, db_alias, shard_id):
#     cursor = connections[db_alias].cursor()
#     cursor.execute(postgres_shard_id_function_sql % {'shard_epoch': settings.SHARD_EPOCH,
#                                                      'shard_id': shard_id,
#                                                      'sequence_name': sequence_name})
#     cursor.close()
# 
# def get_next_sharded_id(shard):
#     cursor = connections[shard].cursor()
#     cursor.execute("SELECT next_sharded_id();")
#     generated_id = cursor.fetchone()
#     cursor.close()
# 
#     return generated_id[0]


from django.apps import apps
from django.conf import settings
from django.db.models import AutoField, CharField, ForeignKey, BigIntegerField, OneToOneField
from django_sharding_library.constants import Backends
from django_sharding_library.utils import create_postgres_global_sequence, create_postgres_shard_id_function, get_next_sharded_id
    from django.db.models import BigAutoField
    from django.utils.translation import ugettext_lazy as _
        from django_sharding_library.id_generation_strategies import TableStrategy
        from django_sharding_library.id_generation_strategies import UUIDStrategy
            from django.db.backends.postgresql.base import DatabaseWrapper as PostgresDatabaseWrapper
            from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as PostgresDatabaseWrapper
            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 BasePostgresShardGeneratedIDField(object):

    def __init__(self, *args, **kwargs):

        if not hasattr(settings, 'SHARD_EPOCH'):
            raise ValueError("PostgresShardGeneratedIDField requires a SHARD_EPOCH to be defined in your settings file.")

        return super(BasePostgresShardGeneratedIDField, self).__init__(*args, **kwargs)

    @staticmethod
    def migration_receiver(*args, **kwargs):
        sequence_name = "global_id_sequence"
        db_alias = kwargs.get('using')
        if not db_alias:
            raise EnvironmentError("A pre-migration receiver did not receive a database alias. "
                                   "Perhaps your app is not registered correctly?")

# Path: django_sharding_library/constants.py
# 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')
# 
# Path: django_sharding_library/utils.py
# def create_postgres_global_sequence(sequence_name, db_alias, reset_sequence=False):
#     cursor = connections[db_alias].cursor()
#     sid = transaction.savepoint(db_alias)
#     create_sequence_if_not_exists_sql = """DO
# $$
# BEGIN
#         CREATE SEQUENCE %s;
# EXCEPTION WHEN duplicate_table THEN
#         -- do nothing, it's already there
# END
# $$ LANGUAGE plpgsql;"""
#     try:
#         cursor.execute(create_sequence_if_not_exists_sql % sequence_name)
#     except DatabaseError:
#         transaction.savepoint_rollback(sid, using=db_alias)
#     else:
#         transaction.savepoint_commit(sid, using=db_alias)
#     if reset_sequence:
#         cursor.execute("SELECT setval('%s', 1, false)" % (sequence_name,))
#     cursor.close()
# 
# def create_postgres_shard_id_function(sequence_name, db_alias, shard_id):
#     cursor = connections[db_alias].cursor()
#     cursor.execute(postgres_shard_id_function_sql % {'shard_epoch': settings.SHARD_EPOCH,
#                                                      'shard_id': shard_id,
#                                                      'sequence_name': sequence_name})
#     cursor.close()
# 
# def get_next_sharded_id(shard):
#     cursor = connections[shard].cursor()
#     cursor.execute("SELECT next_sharded_id();")
#     generated_id = cursor.fetchone()
#     cursor.close()
# 
#     return generated_id[0]


# Path: django_sharding_library/fields.py
from django.apps import apps
from django.conf import settings
from django.db.models import AutoField, CharField, ForeignKey, BigIntegerField, OneToOneField
from django_sharding_library.constants import Backends
from django_sharding_library.utils import create_postgres_global_sequence, create_postgres_shard_id_function, get_next_sharded_id
    from django.db.models import BigAutoField
    from django.utils.translation import ugettext_lazy as _
        from django_sharding_library.id_generation_strategies import TableStrategy
        from django_sharding_library.id_generation_strategies import UUIDStrategy
            from django.db.backends.postgresql.base import DatabaseWrapper as PostgresDatabaseWrapper
            from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as PostgresDatabaseWrapper

            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 BasePostgresShardGeneratedIDField(object):

    def __init__(self, *args, **kwargs):

        if not hasattr(settings, 'SHARD_EPOCH'):
            raise ValueError("PostgresShardGeneratedIDField requires a SHARD_EPOCH to be defined in your settings file.")

        return super(BasePostgresShardGeneratedIDField, self).__init__(*args, **kwargs)

    @staticmethod
    def migration_receiver(*args, **kwargs):
        sequence_name = "global_id_sequence"
        db_alias = kwargs.get('using')
        if not db_alias:
            raise EnvironmentError("A pre-migration receiver did not receive a database alias. "
                                   "Perhaps your app is not registered correctly?")

        if settings.DATABASES[db_alias]['ENGINE'] in Backends.POSTGRES: