Coverage for oc_ocdm / counter_handler / redis_counter_handler.py: 75%
68 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-07-06 20:05 +0000
1#!/usr/bin/python
3# SPDX-FileCopyrightText: 2024-2025 Arcangelo Massari <arcangelo.massari@unibo.it>
4#
5# SPDX-License-Identifier: ISC
7# -*- coding: utf-8 -*-
9from collections.abc import Callable
10from typing import Dict, Optional, Tuple, Union, cast
12import redis
13from redis.client import Pipeline
14from tqdm import tqdm
16from oc_ocdm.counter_handler.counter_handler import CounterHandler
19class RedisCounterHandler(CounterHandler):
20 """A concrete implementation of the ``CounterHandler`` interface that persistently stores
21 the counter values within a Redis database."""
23 def __init__(self, host: str = "localhost", port: int = 6379, db: int = 0, password: Optional[str] = None) -> None:
24 """
25 Constructor of the ``RedisCounterHandler`` class.
27 :param host: Redis server host
28 :type host: str
29 :param port: Redis server port
30 :type port: int
31 :param db: Redis database number
32 :type db: int
33 :param password: Redis password (if required)
34 :type password: Optional[str]
35 """
36 self.host = host
37 self.port = port
38 self.db = db
39 self.password = password
40 self.redis: redis.Redis = redis.Redis(host=host, port=port, db=db, password=password, decode_responses=True)
42 def __getstate__(self):
43 """Support for pickle serialization."""
44 state = self.__dict__.copy()
45 # Remove Redis connection (not picklable)
46 del state["redis"]
47 return state
49 def __setstate__(self, state: dict[str, object]) -> None:
50 """Support for pickle deserialization."""
51 vars(self).update(state)
52 # Recreate Redis connection
53 self.redis = redis.Redis(
54 host=self.host, port=self.port, db=self.db, password=self.password, decode_responses=True
55 )
57 def set_counter(
58 self,
59 new_value: int,
60 entity_short_name: str,
61 prov_short_name: str = "",
62 identifier: int = 1,
63 supplier_prefix: str = "",
64 ) -> None:
65 """
66 It allows to set the counter value of graph and provenance entities.
68 :param new_value: The new counter value to be set
69 :type new_value: int
70 :param entity_short_name: The short name associated either to the type of the entity itself
71 or, in case of a provenance entity, to the type of the relative graph entity.
72 :type entity_short_name: str
73 :param prov_short_name: In case of a provenance entity, the short name associated to the type
74 of the entity itself. An empty string otherwise.
75 :type prov_short_name: str
76 :param identifier: In case of a provenance entity, the counter value that identifies the relative
77 graph entity. The integer value '1' otherwise.
78 :type identifier: int
79 :param supplier_prefix: The supplier prefix
80 :type supplier_prefix: str
81 :raises ValueError: if ``new_value`` is a negative integer
82 :return: None
83 """
84 if new_value < 0:
85 raise ValueError("new_value must be a non negative integer!")
87 key = self._get_key(entity_short_name, prov_short_name, identifier, supplier_prefix)
88 self.redis.set(key, new_value)
90 def read_counter(
91 self, entity_short_name: str, prov_short_name: str = "", identifier: int = 1, supplier_prefix: str = ""
92 ) -> int:
93 """
94 It allows to read the counter value of graph and provenance entities.
96 :param entity_short_name: The short name associated either to the type of the entity itself
97 or, in case of a provenance entity, to the type of the relative graph entity.
98 :type entity_short_name: str
99 :param prov_short_name: In case of a provenance entity, the short name associated to the type
100 of the entity itself. An empty string otherwise.
101 :type prov_short_name: str
102 :param identifier: In case of a provenance entity, the counter value that identifies the relative
103 graph entity. The integer value '1' otherwise.
104 :type identifier: int
105 :param supplier_prefix: The supplier prefix
106 :type supplier_prefix: str
107 :return: The requested counter value.
108 """
109 key = self._get_key(entity_short_name, prov_short_name, identifier, supplier_prefix)
110 value = cast(Optional[str], self.redis.get(key))
111 return int(value) if value is not None else 0
113 def increment_counter(
114 self, entity_short_name: str, prov_short_name: str = "", identifier: int = 1, supplier_prefix: str = ""
115 ) -> int:
116 """
117 It allows to increment the counter value of graph and provenance entities by one unit.
119 :param entity_short_name: The short name associated either to the type of the entity itself
120 or, in case of a provenance entity, to the type of the relative graph entity.
121 :type entity_short_name: str
122 :param prov_short_name: In case of a provenance entity, the short name associated to the type
123 of the entity itself. An empty string otherwise.
124 :type prov_short_name: str
125 :param identifier: In case of a provenance entity, the counter value that identifies the relative
126 graph entity. The integer value '1' otherwise.
127 :type identifier: int
128 :param supplier_prefix: The supplier prefix
129 :type supplier_prefix: str
130 :return: The newly-updated (already incremented) counter value.
131 """
132 key = self._get_key(entity_short_name, prov_short_name, identifier, supplier_prefix)
133 return cast(int, self.redis.incr(key))
135 def set_metadata_counter(self, new_value: int, entity_short_name: str, dataset_name: str | None) -> None:
136 """
137 It allows to set the counter value of metadata entities.
139 :param new_value: The new counter value to be set
140 :type new_value: int
141 :param entity_short_name: The short name associated either to the type of the entity itself.
142 :type entity_short_name: str
143 :param dataset_name: In case of a ``Dataset``, its name. Otherwise, the name of the relative dataset.
144 :type dataset_name: str
145 :raises ValueError: if ``new_value`` is a negative integer
146 :return: None
147 """
148 if new_value < 0:
149 raise ValueError("new_value must be a non negative integer!")
151 key = f"metadata:{dataset_name}:{entity_short_name}"
152 self.redis.set(key, new_value)
154 def read_metadata_counter(self, entity_short_name: str, dataset_name: str | None) -> int:
155 """
156 It allows to read the counter value of metadata entities.
158 :param entity_short_name: The short name associated either to the type of the entity itself.
159 :type entity_short_name: str
160 :param dataset_name: In case of a ``Dataset``, its name. Otherwise, the name of the relative dataset.
161 :type dataset_name: str
163 :return: The requested counter value.
164 """
165 key = f"metadata:{dataset_name}:{entity_short_name}"
166 value = cast(Optional[str], self.redis.get(key))
167 return int(value) if value is not None else 0
169 def increment_metadata_counter(self, entity_short_name: str, dataset_name: str | None) -> int:
170 """
171 It allows to increment the counter value of metadata entities by one unit.
173 :param entity_short_name: The short name associated either to the type of the entity itself.
174 :type entity_short_name: str
175 :param dataset_name: In case of a ``Dataset``, its name. Otherwise, the name of the relative dataset.
176 :type dataset_name: str
178 :return: The newly-updated (already incremented) counter value.
179 """
180 key = f"metadata:{dataset_name}:{entity_short_name}"
181 return cast(int, self.redis.incr(key))
183 def _get_key(
184 self,
185 entity_short_name: str,
186 prov_short_name: str = "",
187 identifier: Union[str, int, None] = None,
188 supplier_prefix: str = "",
189 ) -> str:
190 """
191 Generate a Redis key for the given parameters.
193 :param entity_short_name: The short name associated either to the type of the entity itself
194 or, in case of a provenance entity, to the type of the relative graph entity.
195 :type entity_short_name: str
196 :param prov_short_name: In case of a provenance entity, the short name associated to the type
197 of the entity itself. An empty string otherwise.
198 :type prov_short_name: str
199 :param identifier: In case of a provenance entity, the identifier of the relative graph entity.
200 :type identifier: Union[str, int, None]
201 :param supplier_prefix: The supplier prefix
202 :type supplier_prefix: str
203 :return: The generated Redis key
204 :rtype: str
205 """
206 key_parts = [entity_short_name, supplier_prefix]
207 if prov_short_name:
208 key_parts.append(str(identifier))
209 key_parts.append(prov_short_name)
210 return ":".join(filter(None, key_parts))
212 def batch_update_counters(self, updates: Dict[str, Dict[Tuple[str, str], Dict[int, int]]]) -> None:
213 """
214 Perform batch updates of counters, processing 1 million at a time with a progress bar.
216 :param updates: A dictionary structure containing the updates.
217 The structure is as follows:
218 {
219 supplier_prefix: {
220 (short_name, prov_short_name): {
221 identifier: counter_value
222 }
223 }
224 }
225 :type updates: Dict[str, Dict[Tuple[str, str], Dict[int, int]]]
226 """
227 all_updates: list[tuple[str, int]] = []
228 for supplier_prefix, value in updates.items():
229 for (short_name, prov_short_name), counters in value.items():
230 for identifier, counter_value in counters.items():
231 key = self._get_key(short_name, prov_short_name, identifier, supplier_prefix)
232 all_updates.append((key, counter_value))
234 total_updates = len(all_updates)
235 batch_size = 1_000_000
237 with tqdm(total=total_updates, desc="Updating counters") as pbar:
238 for i in range(0, total_updates, batch_size):
239 batch = all_updates[i : i + batch_size]
240 pipeline_factory = cast(Callable[[], Pipeline], getattr(self.redis, "pipeline"))
241 pipeline = pipeline_factory()
242 for key, value in batch:
243 pipeline.set(key, value)
244 pipeline.execute()
245 pbar.update(len(batch))