From c821c73a3f894ee0783adf0a7e9fc14817ab1e42 Mon Sep 17 00:00:00 2001
From: Oleg Broytman <phd@phdru.name>
Date: Fri, 1 Nov 2019 00:15:27 +0300
Subject: [PATCH] Feat(sa-merge-all): Use `sa-merge-all.py` to merge dumps; use
 `chdir`

---
 playbooks/debian/sa-merge-all.py  | 70 +++++++++++++++++++++++++++++++
 playbooks/debian/sa-merge-all.yml | 12 ++++--
 2 files changed, 79 insertions(+), 3 deletions(-)
 create mode 100755 playbooks/debian/sa-merge-all.py

diff --git a/playbooks/debian/sa-merge-all.py b/playbooks/debian/sa-merge-all.py
new file mode 100755
index 0000000..c129111
--- /dev/null
+++ b/playbooks/debian/sa-merge-all.py
@@ -0,0 +1,70 @@
+#! /usr/bin/env python
+
+import argparse
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Import')
+    parser.add_argument('-o', '--output', required=True, help='output file')
+    parser.add_argument('input_fnames', nargs='+', help='input files to merge')
+    args = parser.parse_args()
+
+    input_files = {}
+
+    def close_all():
+        for f in input_files.values():
+            f.close()
+
+    for fname in args.input_fnames:
+        try:
+            input_files[fname] = open(fname, 'rtU')
+        except (IOError, OSError):
+            close_all()
+            raise
+
+    # Read and process the first 3 lines
+    counters = {'spam': {}, 'nonspam': {}}
+    for fname, infile in input_files.items():
+        version_line = infile.readline()
+        for _line_no in 1, 2:
+            counter_line = infile.readline().strip()
+            v, counter, tag = counter_line.split('\t')
+            if v != 'v':
+                close_all()
+                raise ValueError(
+                        'Bad v-tag in file %s line %s: unknown v-tag %s, '
+                        'expected "v", got %r' % (fname, counter_line, v))
+            if tag == 'num_spam':
+                try:
+                    counters['spam'][fname] = int(counter)
+                except ValueError:
+                    close_all()
+                    raise
+            elif tag == 'num_nonspam':
+                try:
+                    counters['nonspam'][fname] = int(counter)
+                except ValueError:
+                    close_all()
+                    raise
+            else:
+                close_all()
+                raise ValueError(
+                        'Bad tag in file %s line %s: unknown tag %s, '
+                        'expected "num_spam" or "num_nonspam", got %r' % (
+                            fname, counter_line, tag))
+
+    counters_total = {'spam': 0, 'nonspam': 0}
+    for key, files in counters.items():
+        counters_total[key] = sum(files.values())
+
+    try:
+        with open(args.output, 'wt') as outfile:
+            outfile.write(version_line)
+            for key, value in counters_total.items():
+                outfile.write('v\t%d\tnum_%s\n' % (value, key))
+            for infile in input_files.values():
+                for line in infile:
+                    outfile.write(line)
+    except (IOError, OSError):
+        close_all()
+        raise
+    close_all()
diff --git a/playbooks/debian/sa-merge-all.yml b/playbooks/debian/sa-merge-all.yml
index 3157f30..dbf9b4c 100644
--- a/playbooks/debian/sa-merge-all.yml
+++ b/playbooks/debian/sa-merge-all.yml
@@ -24,7 +24,9 @@
       when: inventory_hostname != 'localhost'
 
     - name: Combine SpamAssassin DB backups
-      shell: "cd ~/tmp && exec cat sa-learn.backup@* >sa-learn.backup"
+      shell: "exec {{ playbook_dir }}/sa-merge-all.py -o sa-learn.backup sa-learn.backup@*"
+      args:
+        chdir: "~/tmp"
       when: inventory_hostname == 'localhost'
 
     - block:
@@ -33,7 +35,9 @@
           src: "~/tmp/sa-learn.backup"
           dest: "~/tmp/sa-learn.backup"
       - name: Restore combined SpamAssassin DB
-        shell: "cd tmp && sa-learn --clear && sa-learn --restore sa-learn.backup && exec sa-learn --sync"
+        shell: "sa-learn --clear && sa-learn --restore sa-learn.backup && exec sa-learn --sync"
+        args:
+          chdir: "~/tmp"
 
       - name: Start SpamAssassin
         become: true
@@ -48,4 +52,6 @@
       when: inventory_hostname != 'localhost'
 
     - name: Cleanup
-      shell: "cd ~/tmp && exec rm sa-learn.backup*"
+      shell: "exec rm sa-learn.backup*"
+      args:
+        chdir: "~/tmp"
-- 
2.39.5