Support multiple databases (refs #9)

Author: copelco

.gitignore

@@ -3,6 +3,8 @@
 # Virtualenvs for testing
 .venv
 venv/
+.direnv
+.envrc
 
 # Temporary testing directories
 test__*

README.md

@@ -382,6 +382,25 @@ Provides args:
 * `sender` - Always `None`
 
 
+### Multiple databases
+
+django-pgviews can use multiple databases.  Similar to Django's `migrate`
+management command, our commands (`clear_pgviews`, `refresh_pgviews`,
+`sync_pgviews`) operate on one database at a time. You can specify which
+database to synchronize by providing the `--database` option. For example:
+
+```shell
+python manage.py sync_pgviews  # uses default db
+python manage.py sync_pgviews --database=myotherdb
+```
+
+Unless using custom routers, django-pgviews will sync all views to the specified
+database. If you want to interact with multiple databases automatically, you'll
+need to take some additional steps. Please refer to Django's [Automatic database
+routing](https://docs.djangoproject.com/en/3.2/topics/db/multi-db/#automatic-database-routing)
+to pin views to specific databases.
+
+
 ## Django Compatibility
 
 <table>

django_pgviews/apps.py

@@ -16,7 +16,7 @@ class ViewConfig(apps.AppConfig):
     name = "django_pgviews"
     verbose_name = "Django Postgres Views"
 
-    def sync_pgviews(self, sender, app_config, **kwargs):
+    def sync_pgviews(self, sender, app_config, using, **kwargs):
         """
         Forcibly sync the views.
         """
@@ -35,6 +35,7 @@ def sync_pgviews(self, sender, app_config, **kwargs):
                 force=True,
                 update=True,
                 materialized_views_check_sql_changed=getattr(settings, "MATERIALIZED_VIEWS_CHECK_SQL_CHANGED", False),
+                using=using,
             )
             self.counter = 0
 

django_pgviews/management/commands/clear_pgviews.py

@@ -2,7 +2,7 @@
 
 from django.apps import apps
 from django.core.management.base import BaseCommand
-from django.db import connection
+from django.db import DEFAULT_DB_ALIAS
 
 from django_pgviews.view import clear_view, View, MaterializedView
 
@@ -12,11 +12,21 @@
 class Command(BaseCommand):
     help = """Clear Postgres views. Use this before running a migration"""
 
-    def handle(self, **options):
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--database",
+            default=DEFAULT_DB_ALIAS,
+            help='Nominates a database to synchronize. Defaults to the "default" database.',
+        )
+
+    def handle(self, database, **options):
         for view_cls in apps.get_models():
             if not (isinstance(view_cls, type) and issubclass(view_cls, View) and hasattr(view_cls, "sql")):
                 continue
             python_name = "{}.{}".format(view_cls._meta.app_label, view_cls.__name__)
+            connection = view_cls.get_view_connection(using=database)
+            if not connection:
+                continue
             status = clear_view(
                 connection, view_cls._meta.db_table, materialized=isinstance(view_cls(), MaterializedView)
             )

django_pgviews/management/commands/refresh_pgviews.py

@@ -1,4 +1,5 @@
 from django.core.management.base import BaseCommand
+from django.db import DEFAULT_DB_ALIAS
 
 from django_pgviews.models import ViewRefresher
 
@@ -14,6 +15,11 @@ def add_arguments(self, parser):
             dest="concurrently",
             help="Refresh concurrently if the materialized view supports it",
         )
+        parser.add_argument(
+            "--database",
+            default=DEFAULT_DB_ALIAS,
+            help='Nominates a database to synchronize. Defaults to the "default" database.',
+        )
 
-    def handle(self, concurrently, **options):
-        ViewRefresher().run(concurrently)
+    def handle(self, concurrently, database, **options):
+        ViewRefresher().run(concurrently, using=database)

django_pgviews/management/commands/sync_pgviews.py

@@ -1,5 +1,6 @@
 from django.conf import settings
 from django.core.management.base import BaseCommand
+from django.db import DEFAULT_DB_ALIAS
 
 from django_pgviews.models import ViewSyncer
 
@@ -46,11 +47,16 @@ def add_arguments(self, parser):
                 "if the SQL is different. By default uses django setting MATERIALIZED_VIEWS_CHECK_SQL_CHANGED."
             ),
         )
+        parser.add_argument(
+            "--database",
+            default=DEFAULT_DB_ALIAS,
+            help='Nominates a database to synchronize. Defaults to the "default" database.',
+        )
 
-    def handle(self, force, update, materialized_views_check_sql_changed, **options):
+    def handle(self, force, update, materialized_views_check_sql_changed, database, **options):
         vs = ViewSyncer()
 
         if materialized_views_check_sql_changed is None:
             materialized_views_check_sql_changed = getattr(settings, "MATERIALIZED_VIEWS_CHECK_SQL_CHANGED", False)
 
-        vs.run(force, update, materialized_views_check_sql_changed=materialized_views_check_sql_changed)
+        vs.run(force, update, using=database, materialized_views_check_sql_changed=materialized_views_check_sql_changed)

django_pgviews/models.py

@@ -1,7 +1,6 @@
 import logging
 
 from django.apps import apps
-from django.db import connection
 
 from django_pgviews.signals import view_synced, all_views_synced
 from django_pgviews.view import create_view, View, MaterializedView, create_materialized_view
@@ -38,13 +37,16 @@ def run_backlog(self, backlog, **kwargs):
 
 
 class ViewSyncer(RunBacklog):
-    def run(self, force, update, materialized_views_check_sql_changed=False, **options):
+    def run(self, force, update, using, materialized_views_check_sql_changed=False, **options):
         if super().run(
-            force=force, update=update, materialized_views_check_sql_changed=materialized_views_check_sql_changed
+            force=force,
+            update=update,
+            using=using,
+            materialized_views_check_sql_changed=materialized_views_check_sql_changed,
         ):
-            all_views_synced.send(sender=None)
+            all_views_synced.send(sender=None, using=using)
 
-    def run_backlog(self, backlog, *, force, update, materialized_views_check_sql_changed, **kwargs):
+    def run_backlog(self, backlog, *, force, update, using, materialized_views_check_sql_changed, **kwargs):
         """Installs the list of models given from the previous backlog
 
         If the correct dependent views have not been installed, the view
@@ -67,6 +69,10 @@ def run_backlog(self, backlog, *, force, update, materialized_views_check_sql_ch
                 continue  # Skip
 
             try:
+                connection = view_cls.get_view_connection(using=using)
+                if not connection:
+                    logger.info("Skipping pgview %s (migrations not allowed on %s)", name, using)
+                    continue  # Skip
                 if isinstance(view_cls(), MaterializedView):
                     status = create_materialized_view(
                         connection, view_cls, check_sql_changed=materialized_views_check_sql_changed
@@ -86,6 +92,7 @@ def run_backlog(self, backlog, *, force, update, materialized_views_check_sql_ch
                     force=force,
                     status=status,
                     has_changed=status not in ("EXISTS", "FORCE_REQUIRED"),
+                    using=using,
                 )
                 self.finished.append(name)
             except Exception as exc:
@@ -114,10 +121,10 @@ def run_backlog(self, backlog, *, force, update, materialized_views_check_sql_ch
 
 
 class ViewRefresher(RunBacklog):
-    def run(self, concurrently, **kwargs):
-        return super().run(concurrently=concurrently, **kwargs)
+    def run(self, concurrently, using, **kwargs):
+        return super().run(concurrently=concurrently, using=using, **kwargs)
 
-    def run_backlog(self, backlog, *, concurrently, **kwargs):
+    def run_backlog(self, backlog, *, concurrently, using, **kwargs):
         new_backlog = []
         for view_cls in backlog:
             skip = False
@@ -132,6 +139,11 @@ def run_backlog(self, backlog, *, concurrently, **kwargs):
                 logger.info("Putting pgview at back of queue: %s", name)
                 continue  # Skip
 
+            # Don't refresh views not associated with this database
+            connection = view_cls.get_view_connection(using=using)
+            if not connection:
+                continue
+
             if issubclass(view_cls, MaterializedView):
                 view_cls.refresh(concurrently=concurrently)
                 logger.info("pgview %s refreshed", name)

django_pgviews/signals.py

@@ -1,5 +1,5 @@
 from django.dispatch import Signal
 
 
-view_synced = Signal(providing_args=["update", "force", "status", "has_changed"])
-all_views_synced = Signal()
+view_synced = Signal(providing_args=["update", "force", "status", "has_changed", "using"])
+all_views_synced = Signal(providing_args=["using"])

django_pgviews/view.py

@@ -9,7 +9,7 @@
 import psycopg2
 from django.apps import apps
 from django.core import exceptions
-from django.db import connection, transaction
+from django.db import connections, router, transaction
 from django.db import models
 from django.db.models.query import QuerySet
 
@@ -371,6 +371,18 @@ class View(models.Model, metaclass=ViewMeta):
     def get_sql(cls):
         return ViewSQL(cls.sql, None)
 
+    @classmethod
+    def get_view_connection(cls, using):
+        """
+        Returns connection for "using" database if migrations are allowed (via
+        router). Returns None if migrations are not allowed to indicate view
+        should not be used on the specified database.
+
+        Overwrite this method in subclass to customize, if needed.
+        """
+        if router.allow_migrate(using, cls._meta.app_label):
+            return connections[using]
+
     class Meta:
         abstract = True
         managed = False
@@ -438,7 +450,11 @@ class MaterializedView(View):
 
     @classmethod
     def refresh(self, concurrently=False):
-        cursor_wrapper = connection.cursor()
+        conn = self.get_view_connection(using=router.db_for_write(self))
+        if not conn:
+            logger.warning("Failed to find connection to refresh %s", self)
+            return
+        cursor_wrapper = conn.cursor()
         cursor = cursor_wrapper.cursor
         try:
             if self._concurrent_index is not None and concurrently: