def test_with_context(self):
with ConsulServiceController().start_service() as service:
consul_client = service.create_consul_client()
lock_manager = ConsulLockManager(consul_client=consul_client)
with lock_manager.acquire(KEY_1) as lock_information:
self.assertEqual(lock_information, lock_manager.find(KEY_1))
self.assertIsNone(lock_manager.find(KEY_1))
def _acquire_lock(lock_manager: ConsulLockManager, configuration: CliLockConfiguration) \
-> Optional[ConnectedConsulLockInformation]:
"""
TODO
:param lock_manager:
:param configuration:
:return:
"""
event_listeners: LockEventListener = {}
if configuration.on_before_locked_executables is not None:
event_listeners["on_before_lock"] = _generate_event_listener_caller(
configuration.on_before_locked_executables)
if configuration.on_lock_already_locked_executables is not None:
event_listeners["on_lock_already_locked"] = _generate_event_listener_caller(
configuration.on_lock_already_locked_executables)
try:
return lock_manager.acquire(
key=configuration.key, blocking=not configuration.non_blocking,
timeout=configuration.timeout, metadata=configuration.metadata, **event_listeners,
lock_poll_interval_generator=lambda i: configuration.lock_poll_interval)
except LockAcquireTimeoutError as e:
logger.debug(e)
logger.error(f"Timed out whilst waiting to acquire lock: {configuration.key}")
print(json.dumps(None))
exit(LOCK_ACQUIRE_TIMEOUT_EXIT_CODE)
def test_with_context(self):
with ConsulServiceController().start_service() as service:
consul_client = service.create_consul_client()
lock_manager = ConsulLockManager(consul_client=consul_client)
with lock_manager.acquire(KEY_1) as lock_information:
self.assertEqual(lock_information, lock_manager.find(KEY_1))
self.assertIsNone(lock_manager.find(KEY_1))
def test_unlock_all(self):
test_keys = [f"{KEY_1}_{i}" for i in range(5)]
with ConsulServiceController().start_service() as service:
lock_manager = ConsulLockManager(consul_client=service.create_consul_client())
for key in test_keys:
lock = lock_manager.acquire(key)
assert isinstance(lock, ConsulLockInformation)
unlock_results = lock_manager.release_all(test_keys)
for unlock_result in unlock_results:
self.assertTrue(unlock_result)
def test_unlock_all(self):
test_keys = [f"{KEY_1}_{i}" for i in range(5)]
with ConsulServiceController().start_service() as service:
lock_manager = ConsulLockManager(
consul_client=service.create_consul_client())
for key in test_keys:
lock = lock_manager.acquire(key)
assert isinstance(lock, ConsulLockInformation)
unlock_results = lock_manager.release_all(test_keys)
for unlock_result in unlock_results:
self.assertTrue(unlock_result)
def first_locker(key: str, service: ConsulDockerisedService) -> CaptureResult:
lock_manager = ConsulLockManager(consul_client=service.create_consul_client())
lock_information = lock_manager.acquire(KEY_1)
return CaptureResult(return_value=lock_information)
class ConsulChecksumStorage(ChecksumStorage):
"""
Consul storage for configuration -> checksum mappings.
"""
CONSUL_HTTP_TOKEN_ENVIRONMENT_VARIABLE = "CONSUL_HTTP_TOKEN"
CONSUL_SESSION_LOCK_DEFAULT_TIMEOUT = 120
TEXT_ENCODING = "utf-8"
_IMPORT_MISSING_ERROR_MESSAGE = "To use Consul storage, please install the requirements in " \
"`consul_requirements.txt`"
@staticmethod
def _load_consul_class() -> Type:
"""
Loads the Consul class at run time (optional requirement).
:return: the Consul class
:raises MissingOptionalDependencyError: if a required dependency is not installed
"""
try:
from consul import Consul
except ImportError as e:
raise MissingOptionalDependencyError(
ConsulChecksumStorage._IMPORT_MISSING_ERROR_MESSAGE) from e
return Consul
@staticmethod
def _load_consul_lock_manager() -> Type:
"""
Loads the ConsulLockManager class at run time (optional requirement).
:return: the Consul class
:raises MissingOptionalDependencyError: if a required dependency is not installed
"""
try:
from consullock.managers import ConsulLockManager
except ImportError as e:
raise MissingOptionalDependencyError(
ConsulChecksumStorage._IMPORT_MISSING_ERROR_MESSAGE) from e
return ConsulLockManager
@property
def url(self) -> str:
return self._consul_client.http.base_uri
@property
def token(self) -> str:
return self._consul_client.token
def __init__(self,
data_key: str,
lock_key: str,
url: str = None,
token: str = None,
consul_client=None,
configuration_checksum_mappings: Mapping[str, str] = None):
Consul = ConsulChecksumStorage._load_consul_class()
ConsulLockManager = ConsulChecksumStorage._load_consul_lock_manager()
if url is not None and consul_client is not None:
raise ValueError("Cannot use both `url` and `consul_client`")
self.data_key = data_key
self.lock_key = lock_key
consul_client_kwargs: Dict = {}
if url is not None:
parsed_url = urlparse(url)
consul_client_kwargs["host"] = parsed_url.hostname
consul_client_kwargs["port"] = parsed_url.port
consul_client_kwargs["scheme"] = parsed_url.scheme if len(
parsed_url.scheme) > 0 else "http"
self._consul_client = consul_client if consul_client is not None else Consul(
**consul_client_kwargs)
if token is None:
token = os.environ.get(
ConsulChecksumStorage.CONSUL_HTTP_TOKEN_ENVIRONMENT_VARIABLE,
None)
if token is not None:
# Work around for https://github.com/cablehead/python-consul/issues/170
self._consul_client.token = token
self._consul_client.http.session.headers.update(
{"X-Consul-Token": token})
self._lock_manager = ConsulLockManager(
consul_client=self._consul_client,
session_ttl_in_seconds=ConsulChecksumStorage.
CONSUL_SESSION_LOCK_DEFAULT_TIMEOUT)
super().__init__(configuration_checksum_mappings)
def get_checksum(self, configuration_id: str) -> Optional[str]:
return self.get_all_checksums().get(configuration_id)
def get_all_checksums(self) -> Dict[str, str]:
value = self._consul_client.kv.get(self.data_key)[1]
if value is None:
return {}
value = value["Value"].decode(ConsulChecksumStorage.TEXT_ENCODING)
return json.loads(value)
def set_checksum(self, configuration_id: str, checksum: str):
with self._lock_manager.acquire(self.lock_key):
value = self.get_all_checksums()
value[configuration_id] = checksum
self._consul_client.kv.put(self.data_key,
json.dumps(value, sort_keys=True))
def set_all_checksums(self, configuration_checksum_mappings: Mapping[str,
str]):
with self._lock_manager.acquire(self.lock_key):
value = self.get_all_checksums()
value.update(configuration_checksum_mappings)
self._consul_client.kv.put(self.data_key,
json.dumps(value, sort_keys=True))
def first_locker(key: str,
service: ConsulDockerisedService) -> CaptureResult:
lock_manager = ConsulLockManager(
consul_client=service.create_consul_client())
lock_information = lock_manager.acquire(KEY_1)
return CaptureResult(return_value=lock_information)
