make psycog2 optional by pulling down the import (#63)

JBKahn · Mar 19, 2017 · 2cc8532 · 2cc8532
1 parent 58d633b
commit 2cc8532
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/django_sharding_library/fields.py b/django_sharding_library/fields.py
@@ -5,10 +5,6 @@
 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
 
-try:
-    from django.db.backends.postgresql.base import DatabaseWrapper as PostgresDatabaseWrapper
-except ImportError:
-    from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as PostgresDatabaseWrapper
 
 try:
     from django.db.models import BigAutoField
@@ -196,6 +192,11 @@ class PostgresShardGeneratedIDAutoField(BasePostgresShardGeneratedIDField, BigAu
     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: