Ajout d'une sous-commande de nettoyage (clean)

Nettoyage des tables de travail utilisées lors de l'import

dora/sirene/backup.py

@@ -73,14 +73,29 @@ def _suffix():
 
 def rename_table(orig_table_name: str, dest_table_name: str):
     with connection.cursor() as c:
-        c.execute("ALTER TABLE %s RENAME TO %s;", [orig_table_name, dest_table_name])
+        # au lieu de :
+        # c.execute("ALTER TABLE %s RENAME TO %s;", [orig_table_name, dest_table_name])
+        # les placeholders de type `string` sont automatiquement entourés de quotes ('),
+        # ce qui vrille les ordres DDL.
+        order = f"ALTER TABLE {orig_table_name} RENAME TO {dest_table_name};"
+        c.execute(order)
 
 
 def vacuum_analyze():
     with connection.cursor() as c:
         c.execute("VACUUM ANALYZE;")
 
 
+def clean_tmp_tables(*tmp_tables):
+    # attention ...
+    with connection.cursor() as c:
+        # petite sécurité supplémentaire:
+        # les tables temporaires sont obligatoirement préfixées par `_`
+        for tmp_table in [tt for tt in tmp_tables if tt.startswith("_")]:
+            print(f" > suppression de {tmp_table}")
+            c.execute(f"DROP TABLE IF EXISTS {tmp_table}")
+
+
 def create_insert_statement(table_name: str) -> tuple[str, list[str]]:
     fields = [f.name for f in Establishment._meta.fields]
     stmt = f"INSERT INTO public.{table_name}({",".join(fields)}) VALUES({ ",".join(["%s"]*len(fields)) })"

dora/sirene/management/commands/import_sirene.py

@@ -11,6 +11,7 @@
 
 from dora.sirene.backup import (
     bulk_add_establishments,
+    clean_tmp_tables,
     create_indexes,
     create_table,
     rename_table,
@@ -55,6 +56,12 @@ def add_arguments(self, parser):
             help="Effectue un VACUUM ANALYZE sur la base.",
         )
 
+        parser.add_argument(
+            "--clean",
+            action="store_true",
+            help="Efface les tables de travail temporaires en DB.",
+        )
+
     def download_data(self, tmp_dir_name):
         if USE_TEMP_DIR:
             the_dir = pathlib.Path(tmp_dir_name)
@@ -68,7 +75,11 @@ def download_data(self, tmp_dir_name):
         zipped_stock_file = the_dir / "StockUniteLegale_utf8.zip"
 
         if not os.path.exists(zipped_stock_file):
-            self.stdout.write(self.style.NOTICE("Téléchargement des 'unités légales' (entreprises mères)"))
+            self.stdout.write(
+                self.style.NOTICE(
+                    "Téléchargement des 'unités légales' (entreprises mères)"
+                )
+            )
             subprocess.run(
                 ["curl", legal_units_file_url, "-o", zipped_stock_file],
                 check=True,
@@ -92,7 +103,9 @@ def download_data(self, tmp_dir_name):
                 check=True,
             )
 
-            self.stdout.write(self.style.NOTICE("Décompression du fichier établissements"))
+            self.stdout.write(
+                self.style.NOTICE("Décompression du fichier établissements")
+            )
             subprocess.run(
                 ["gzip", "-dk", gzipped_estab_file],
                 check=True,
@@ -165,33 +178,42 @@ def handle(self, *args, **options):
         if options.get("activate"):
             # activation de la table temporaire (si existante),
             # comme table de production (`sirene_establishment`)
-            self.stdout.write(self.WARNING("Activation de la table de travail"))
+            self.stdout.write(self.style.WARNING("Activation de la table de travail"))
 
             # on sauvegarde la base de production
-            self.stdout.write(self.NOTICE(" > sauvegarde de la table actuelle"))
+            self.stdout.write(self.style.NOTICE(" > sauvegarde de la table actuelle"))
             rename_table(SIRENE_TABLE, BACKUP_TABLE)
 
             # on renomme la table de travail
-            self.stdout.write(self.NOTICE(" > renommage de la table de travail"))
+            self.stdout.write(self.style.NOTICE(" > renommage de la table de travail"))
             rename_table(TMP_TABLE, SIRENE_TABLE)
 
-            self.stdout.write(self.NOTICE("Activation terminée"))
+            self.stdout.write(self.style.NOTICE("Activation terminée"))
             return
 
         if options.get("rollback"):
             # activation de la table sauvegardée
-            self.stdout.write(self.WARNING("Activation de la table sauvegardée"))
+            self.stdout.write(self.style.WARNING("Activation de la table sauvegardée"))
             rename_table(SIRENE_TABLE, TMP_TABLE)
             rename_table(BACKUP_TABLE, SIRENE_TABLE)
             rename_table(TMP_TABLE, BACKUP_TABLE)
 
-        if options.get("analyse"):
+        if options.get("analyze"):
             # lance une analyse statistique sur la base Postgres
             self.stdout.write(self.style.WARNING("Analyse de la DB en cours..."))
             vacuum_analyze()
             self.stdout.write(self.style.NOTICE("Analyse terminée"))
             return
 
+        if options.get("clean"):
+            # Supprime les tables de travail / temporaires de la base Postgres
+            self.stdout.write(
+                self.style.WARNING("Suppression des tables temporaires...")
+            )
+            clean_tmp_tables(TMP_TABLE, BACKUP_TABLE)
+            self.stdout.write(self.style.NOTICE("Suppression terminée"))
+            return
+
         self.stdout.write(self.style.NOTICE(" > création de la base de travail"))
         # efface la précédente
         create_table(TMP_TABLE)
@@ -277,4 +299,8 @@ def handle(self, *args, **options):
                 # la sauvegarde de la base de production et l'analyse de la DB
                 # ne sont pas automatique, voir arguments `--activate` et `--analyze`
 
-                self.stdout.write(self.style.SUCCESS("L'import est terminé. Ne pas oublier d'activer la table de travail (--activate)"))
+                self.stdout.write(
+                    self.style.SUCCESS(
+                        "L'import est terminé. Ne pas oublier d'activer la table de travail (--activate)"
+                    )
+                )