fix dynamic creation of measure columns in the correct database

frePPLe · Feb 9, 2025 · 28fb81d · 28fb81d
1 parent 6375e88
commit 28fb81d
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 37 deletions.
diff --git a/freppledb/execute/management/commands/importworkbook.py b/freppledb/execute/management/commands/importworkbook.py
@@ -316,7 +316,7 @@ def handle(self, **options):
                 if "freppledb.forecast" in settings.INSTALLED_APPS:
                     from freppledb.forecast.models import ForecastPlan
 
-                    ForecastPlan.refreshTableColumns()
+                    ForecastPlan.refreshTableColumns(self.database)
 
             except GeneratorExit:
                 logger.warning("Connection Aborted")

diff --git a/freppledb/execute/management/commands/loaddata.py b/freppledb/execute/management/commands/loaddata.py
@@ -140,7 +140,7 @@ def handle(self, *fixture_labels, **options):
             if "freppledb.forecast" in settings.INSTALLED_APPS:
                 from freppledb.forecast.models import ForecastPlan
 
-                ForecastPlan.refreshTableColumns()
+                ForecastPlan.refreshTableColumns(database)
 
             # if the fixture doesn't contain the 'demo' word, let's not apply loaddata post-treatments
             if "FREPPLE_TEST" in os.environ:

diff --git a/freppledb/forecast/commands.py b/freppledb/forecast/commands.py
@@ -642,8 +642,10 @@ def run(cls, database=DEFAULT_DB_ALIAS, **kwargs):
             ]
 
             if not all(m in measures for m in mandatory):
-                management.call_command("loaddata", "measures.json", verbosity=0)
-                ForecastPlan.refreshTableColumns()
+                management.call_command(
+                    "loaddata", "measures.json", database=database, verbosity=0
+                )
+                ForecastPlan.refreshTableColumns(database)
                 logger.info("Some missing mandatory measures have been added.")
 
 

diff --git a/freppledb/forecast/models.py b/freppledb/forecast/models.py
@@ -898,7 +898,7 @@ def sendToService(d):
             cursor.execute("REFRESH MATERIALIZED VIEW forecastreport_view")
 
     @staticmethod
-    def refreshTableColumns():
+    def refreshTableColumns(database=DEFAULT_DB_ALIAS):
         """
         Adjust the forecastplan table to have a columnn for every measure.
 
@@ -909,39 +909,36 @@ def refreshTableColumns():
         expected_columns = [
             m.name for m in Measure.standard_measures() if not m.computed
         ]
-        for m in Measure.objects.using(DEFAULT_DB_ALIAS).only("name"):
+        for m in Measure.objects.using(database).only("name"):
             expected_columns.append(m.name)
 
-        # Loop over the in-use scenarios, and check the forecastplan table
+        # Check the forecastplan table
         modified = False
-        for db in Scenario.objects.using(DEFAULT_DB_ALIAS).filter(
-            models.Q(status="In use") | models.Q(name=DEFAULT_DB_ALIAS)
-        ):
-            with connections[db.name].cursor() as cursor:
-                cursor.execute(
-                    """
-                    select column_name
-                    from information_schema.columns 
-                    where table_name = 'forecastplan'
-                    and column_name not in (
-                      'item_id', 'location_id', 'customer_id', 'startdate', 'enddate'
-                      )
-                    """
-                )
-                columns = [c[0] for c in cursor.fetchall()]
-                for m in expected_columns:
-                    if m not in columns:
-                        cursor.execute(
-                            "alter table forecastplan add column if not exists %s decimal(20,8)"
-                            % m
-                        )
-                        modified = True
-                for c in columns:
-                    if c not in expected_columns:
-                        cursor.execute(
-                            "alter table forecastplan drop column if exists %s" % c
-                        )
-                        modified = True
+        with connections[database].cursor() as cursor:
+            cursor.execute(
+                """
+                select column_name
+                from information_schema.columns 
+                where table_name = 'forecastplan'
+                and column_name not in (
+                    'item_id', 'location_id', 'customer_id', 'startdate', 'enddate'
+                    )
+                """
+            )
+            columns = [c[0] for c in cursor.fetchall()]
+            for m in expected_columns:
+                if m not in columns:
+                    cursor.execute(
+                        "alter table forecastplan add column if not exists %s decimal(20,8)"
+                        % m
+                    )
+                    modified = True
+            for c in columns:
+                if c not in expected_columns:
+                    cursor.execute(
+                        "alter table forecastplan drop column if exists %s" % c
+                    )
+                    modified = True
 
         if modified:
             # Trigger wsgi reload by importing from freppledb.common.utils.force
@@ -1395,14 +1392,14 @@ def save(self, *args, **kwargs):
         super().save(*args, **kwargs)
 
         # Add or update the database schema
-        ForecastPlan.refreshTableColumns()
+        ForecastPlan.refreshTableColumns(database=self._state.db)
 
     def delete(self, *args, **kwargs):
         # Call the real save() method
         super().delete(*args, **kwargs)
 
         # Add or update the database schema
-        ForecastPlan.refreshTableColumns()
+        ForecastPlan.refreshTableColumns(database=self._state.db)
 
     class Meta(AuditModel.Meta):
         db_table = "measure"