Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/c7n/cache.py: 39%

1# Copyright The Cloud Custodian Authors. 

2# SPDX-License-Identifier: Apache-2.0 

3"""Provide basic caching services to avoid extraneous queries over 

4multiple policies on the same resource type. 

5""" 

6import pickle # nosec nosemgrep 

7 

8from datetime import datetime, timedelta 

9import os 

10import logging 

11import sqlite3 

12 

13log = logging.getLogger('custodian.cache') 

14 

15CACHE_NOTIFY = False 

16 

17 

18def factory(config): 

19 

20 global CACHE_NOTIFY 

21 

22 if not config: 

23 return NullCache(None) 

24 

25 if not config.cache or not config.cache_period: 

26 if not CACHE_NOTIFY: 

27 log.debug("Disabling cache") 

28 CACHE_NOTIFY = True 

29 return NullCache(config) 

30 elif config.cache == 'memory': 

31 if not CACHE_NOTIFY: 

32 log.debug("Using in-memory cache") 

33 CACHE_NOTIFY = True 

34 return InMemoryCache(config) 

35 return SqlKvCache(config) 

36 

37 

38class Cache: 

39 

40 def __init__(self, config): 

41 self.config = config 

42 

43 def load(self): 

44 return False 

45 

46 def get(self, key): 

47 pass 

48 

49 def save(self, key, data): 

50 pass 

51 

52 def size(self): 

53 return 0 

54 

55 def close(self): 

56 pass 

57 

58 def __enter__(self): 

59 self.load() 

60 return self 

61 

62 def __exit__(self, exc_type, exc_val, exc_tb): 

63 self.close() 

64 

65 

66class NullCache(Cache): 

67 pass 

68 

69 

70class InMemoryCache(Cache): 

71 # Running in a temporary environment, so keep as a cache. 

72 

73 __shared_state = {} 

74 

75 def __init__(self, config): 

76 super().__init__(config) 

77 self.data = self.__shared_state 

78 

79 def load(self): 

80 return True 

81 

82 def get(self, key): 

83 return self.data.get(encode(key)) 

84 

85 def save(self, key, data): 

86 self.data[encode(key)] = data 

87 

88 def size(self): 

89 return sum(map(len, self.data.values())) 

90 

91 

92def encode(key): 

93 return pickle.dumps(key, protocol=pickle.HIGHEST_PROTOCOL) # nosemgrep 

94 

95 

96def resolve_path(path): 

97 return os.path.abspath( 

98 os.path.expanduser( 

99 os.path.expandvars(path))) 

100 

101 

102class SqlKvCache(Cache): 

103 

104 create_table = """ 

105 create table if not exists c7n_cache ( 

106 key blob primary key, 

107 value blob, 

108 create_date timestamp 

109 ) 

110 """ 

111 

112 def __init__(self, config): 

113 super().__init__(config) 

114 self.cache_period = config.cache_period 

115 self.cache_path = resolve_path(config.cache) 

116 self.conn = None 

117 

118 def init(self): 

119 # migration from pickle cache file 

120 if os.path.exists(self.cache_path): 

121 with open(self.cache_path, 'rb') as fh: 

122 header = fh.read(15) 

123 if header != b'SQLite format 3': 

124 log.debug('removing old cache file') 

125 os.remove(self.cache_path) 

126 elif not os.path.exists(os.path.dirname(self.cache_path)): 

127 # parent directory creation 

128 os.makedirs(os.path.dirname(self.cache_path)) 

129 self.conn = sqlite3.connect(self.cache_path) 

130 self.conn.execute(self.create_table) 

131 with self.conn as cursor: 

132 result = cursor.execute( 

133 'delete from c7n_cache where create_date < ?', 

134 [datetime.utcnow() - timedelta(minutes=self.cache_period)]) 

135 if result.rowcount: 

136 log.debug('expired %d stale cache entries', result.rowcount) 

137 

138 def load(self): 

139 if not self.conn: 

140 self.init() 

141 return True 

142 

143 def get(self, key): 

144 with self.conn as cursor: 

145 r = cursor.execute( 

146 'select value, create_date from c7n_cache where key = ?', 

147 [sqlite3.Binary(encode(key))] 

148 ) 

149 row = r.fetchone() 

150 if row is None: 

151 return None 

152 value, create_date = row 

153 create_date = sqlite3.converters['TIMESTAMP'](create_date.encode('utf8')) 

154 if (datetime.utcnow() - create_date).total_seconds() / 60.0 > self.cache_period: 

155 return None 

156 return pickle.loads(value) # nosec nosemgrep 

157 

158 def save(self, key, data, timestamp=None): 

159 with self.conn as cursor: 

160 timestamp = timestamp or datetime.utcnow() 

161 cursor.execute( 

162 'replace into c7n_cache (key, value, create_date) values (?, ?, ?)', 

163 (sqlite3.Binary(encode(key)), sqlite3.Binary(encode(data)), timestamp)) 

164 

165 def size(self): 

166 return os.path.exists(self.cache_path) and os.path.getsize(self.cache_path) or 0 

167 

168 def close(self): 

169 if self.conn: 

170 self.conn.close() 

171 self.conn = None