Handle difference in pickle protocol between python 2.x and 3.x
[imap-fix-internaldate] / src / caching_data.py
1 '''
2 caching_data.py - The module contains the CachingData class.
3
4 Copyright (c) 2012 Intra2net AG
5 Author: Plamen Dimitrov
6
7 This program is free software: you can redistribute it and/or modify
8 it under the terms of the GNU General Public License as published by
9 the Free Software Foundation, either version 3 of the License, or
10 (at your option) any later version.
11
12 This program is distributed in the hope that it will be useful,
13 but WITHOUT ANY WARRANTY; without even the implied warranty of
14 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15 GNU General Public License for more details.
16 '''
17 import os, platform, tempfile
18 import pickle
19 import logging
20 from mailbox_state import MailboxState
21
22 CACHE_FILENAME = "message_cache.dat"
23 CACHE_VERSION = "1"
24
25 class CachingData:
26     """This class is responsible for the caching of data."""
27     
28     # class attributes
29     # integer for version of the cache
30     version = None
31     # boolean flag which indicates fallback mode of the cache
32     fallback_to_date_header = None
33     # dictionary of usernames as keys and dictionaries as values
34     # the second dictionaries have unique mailbox keys and mailboxes as values
35     data = None
36
37     def __init__(self, fallback_mode):
38         # open data file or create one and initialize date if not found
39         try:
40             cachefile = open(CACHE_FILENAME, 'rb')
41             cache_info, self.data = pickle.load(cachefile)
42             cache_info = cache_info.split(' ')
43             self.version = cache_info[0]
44             if(self.version != CACHE_VERSION):
45                 raise IOError("Cache file has version %s and the script version is %s" % (self.version, CACHE_VERSION))
46             self.fallback_to_date_header = cache_info[1]
47             if(self.fallback_to_date_header != str(fallback_mode)):
48                 raise IOError("Cache file date fallback mode setting is different than current settings")
49             logging.info("Cache file %s loaded", CACHE_FILENAME)
50             logging.info("%s users found.", len(self.data))
51         except (IOError, ValueError) as ex:
52             logging.warning("Couldn't load cache file %s: %s", CACHE_FILENAME, ex)
53             logging.warning("DELETING CACHE")
54             self.version = CACHE_VERSION
55             stored_cache_info = self.version + ' ' + str(fallback_mode)
56             self.data = {}
57             with open(CACHE_FILENAME, 'wb') as cachefile:
58                 pickle.dump((stored_cache_info, self.data), cachefile)
59
60     def __del__(self):
61         # create temporary file first
62         location = os.path.dirname(CACHE_FILENAME)    
63         file_descriptor, tmpname = tempfile.mkstemp(dir=location)
64         try:
65             cachefile = os.fdopen(file_descriptor, 'wb')
66
67             # prepare data based on a save flag
68             saved_data = {}
69             for user in self.data:
70                 saved_data[user] = {}
71                 for box_key in self.data[user]:
72                     if(self.data[user][box_key].needs_save):
73                         saved_data[user][box_key] = self.data[user][box_key]
74                         logging.debug("The mailbox %s will be saved.", saved_data[user][box_key].name)
75                 if(len(saved_data[user])==0):
76                     del saved_data[user]
77                     logging.debug("The user %s will not be saved.", user)
78             self.data = saved_data
79             # avoid test mode or cases where nothing needs saving
80             if(len(saved_data)==0):
81                 cachefile.close()
82                 os.unlink(tmpname)
83                 return
84
85             # serialize in file
86             stored_cache_info = self.version + ' ' + self.fallback_to_date_header
87             pickle.dump((stored_cache_info, self.data), cachefile)
88             logging.debug("%s users stored.", len(self.data))
89
90             # handle windows non-atomic rename
91             if(platform.system()=='Windows'):
92                 if(os.path.exists(CACHE_FILENAME)):
93                     cachefile.close()
94                     os.unlink(CACHE_FILENAME)
95
96             os.rename(tmpname, CACHE_FILENAME)
97         except:
98             # clean up temporary file
99             os.unlink(tmpname)
100             raise
101
102         logging.info("Wrote cache file %s", CACHE_FILENAME)
103
104     def retrieve_cached_mailbox(self, name, uidvalidity, user):
105         """Retrieve a cached mailbox or create it."""
106         box_key = name.strip('"') + uidvalidity
107         if(user not in self.data):
108             self.data[user] = {}
109             logging.debug("New user %s cached.", user)
110         if(box_key not in self.data[user]):
111             self.data[user][box_key] = MailboxState(name, uidvalidity, user)
112             logging.debug("New mailbox %s cached.", box_key)
113         return self.data[user][box_key]
114     
115     def report_conflicts(self):
116         """Write a date conflicts report in a file."""
117         with open("conflict_stats.txt", 'w') as statsfile:
118             owner_total_conflicts = {}
119             owner_total_missing = {}
120             for user in self.data:
121                 statsfile.write("user: %s\n" % user)
122                 owner_total_conflicts[user] = 0
123                 owner_total_missing[user] = 0
124                 for box_key in self.data[user]:
125                     owner_total_conflicts[user] += self.data[user][box_key].date_conflicts
126                     owner_total_missing[user] += self.data[user][box_key].no_received_field
127                     statsfile.write("date conflicts: %-15.15s missing header: %-15.15s mailbox: %s\n"\
128                                     % (self.data[user][box_key].date_conflicts,
129                                             self.data[user][box_key].no_received_field,
130                                             self.data[user][box_key].name))
131                 statsfile.write("date conflicts: %-15.15s missing header: %-15.15s TOTAL \n\n"\
132                                 % (owner_total_conflicts[user], owner_total_missing[user]))
133         return