Use configured db schema also for sequences (celery#5972)

the-owl · Feb 27, 2020 · c52105e · c52105e
1 parent 58b8d0d
commit c52105e
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 7 deletions.
diff --git a/celery/backends/database/__init__.py b/celery/backends/database/__init__.py
@@ -89,14 +89,13 @@ def __init__(self, dburi=None, engine_options=None, url=None, **kwargs):
             conf.database_short_lived_sessions)
 
         schemas = conf.database_table_schemas or {}
-        self.task_cls.__table__.schema = schemas.get('task')
-        self.taskset_cls.__table__.schema = schemas.get('group')
-
         tablenames = conf.database_table_names or {}
-        self.task_cls.__table__.name = tablenames.get('task',
-                                                      'celery_taskmeta')
-        self.taskset_cls.__table__.name = tablenames.get('group',
-                                                         'celery_tasksetmeta')
+        self.task_cls.configure(
+            schema=schemas.get('task'),
+            name=tablenames.get('task'))
+        self.taskset_cls.configure(
+            schema=schemas.get('group'),
+            name=tablenames.get('group'))
 
         if not self.url:
             raise ImproperlyConfigured(

diff --git a/celery/backends/database/models.py b/celery/backends/database/models.py
@@ -46,6 +46,12 @@ def to_dict(self):
     def __repr__(self):
         return '<Task {0.task_id} state: {0.status}>'.format(self)
 
+    @classmethod
+    def configure(cls, schema=None, name=None):
+        cls.__table__.schema = schema
+        cls.id.default.schema = schema
+        cls.__table__.name = name or cls.__tablename__
+
 
 class TaskExtended(Task):
     """For the extend result."""
@@ -100,3 +106,9 @@ def to_dict(self):
 
     def __repr__(self):
         return '<TaskSet: {0.taskset_id}>'.format(self)
+
+    @classmethod
+    def configure(cls, schema=None, name=None):
+        cls.__table__.schema = schema
+        cls.id.default.schema = schema
+        cls.__table__.name = name or cls.__tablename__
diff --git a/t/unit/backends/test_database.py b/t/unit/backends/test_database.py
@@ -86,7 +86,18 @@ def test_table_schema_config(self):
         }
         tb = DatabaseBackend(self.uri, app=self.app)
         assert tb.task_cls.__table__.schema == 'foo'
+        assert tb.task_cls.__table__.c.id.default.schema == 'foo'
         assert tb.taskset_cls.__table__.schema == 'bar'
+        assert tb.taskset_cls.__table__.c.id.default.schema == 'bar'
+
+    def test_table_name_config(self):
+        self.app.conf.database_table_names = {
+            'task': 'foo',
+            'group': 'bar',
+        }
+        tb = DatabaseBackend(self.uri, app=self.app)
+        assert tb.task_cls.__table__.name == 'foo'
+        assert tb.taskset_cls.__table__.name == 'bar'
 
     def test_missing_task_id_is_PENDING(self):
         tb = DatabaseBackend(self.uri, app=self.app)