Integrate schema handling within ZKHandler

Abstracts away the schema management, especially when doing actions, to
prevent duplication in other areas.
This commit is contained in:
Joshua Boniface 2021-06-09 13:23:57 -04:00
parent 76c37e6628
commit a9a57533a7
1 changed files with 77 additions and 32 deletions

View File

@ -49,6 +49,7 @@ class ZKConnection(object):
def connection(*args, **kwargs): def connection(*args, **kwargs):
zkhandler = ZKHandler(self.config) zkhandler = ZKHandler(self.config)
zkhandler.connect() zkhandler.connect()
zkhandler.schema.load(zkhandler.read(zkhandler.schema.path('base.schema.version')), quiet=True)
ret = function(zkhandler, *args, **kwargs) ret = function(zkhandler, *args, **kwargs)
@ -87,11 +88,14 @@ class ZKHandler(object):
Initialize an instance of the ZKHandler class with config Initialize an instance of the ZKHandler class with config
A zk_conn object will be created but not started A zk_conn object will be created but not started
A ZKSchema instance will be created
""" """
self.encoding = 'utf8' self.encoding = 'utf8'
self.coordinators = config['coordinators'] self.coordinators = config['coordinators']
self.logger = logger self.logger = logger
self.zk_conn = KazooClient(hosts=self.coordinators) self.zk_conn = KazooClient(hosts=self.coordinators)
self._schema = ZKSchema()
# #
# Class meta-functions # Class meta-functions
@ -105,6 +109,13 @@ class ZKHandler(object):
else: else:
print(message) print(message)
#
# Properties
#
@property
def schema(self):
return self._schema
# #
# State/connection management # State/connection management
# #
@ -130,7 +141,7 @@ class ZKHandler(object):
def connect(self, persistent=False): def connect(self, persistent=False):
""" """
Start the zk_conn object and connect to the cluster Start the zk_conn object and connect to the cluster, then load the current schema version
""" """
try: try:
self.zk_conn.start() self.zk_conn.start()
@ -148,6 +159,23 @@ class ZKHandler(object):
self.zk_conn.stop() self.zk_conn.stop()
self.zk_conn.close() self.zk_conn.close()
#
# Schema helper actions
#
def get_schema_path(self, key):
if isinstance(key, tuple):
# This is a key tuple with both an ipath and an item
ipath, item = key
elif isinstance(key, str):
# This is a key string with just an ipath
ipath = key
item = None
else:
# This is an invalid key
return None
return self.schema.path(ipath, item=item)
# #
# Key Actions # Key Actions
# #
@ -155,7 +183,8 @@ class ZKHandler(object):
""" """
Check if a key exists Check if a key exists
""" """
stat = self.zk_conn.exists(key) path = self.get_schema_path(key)
stat = self.zk_conn.exists(path)
if stat: if stat:
return True return True
else: else:
@ -165,7 +194,8 @@ class ZKHandler(object):
""" """
Read data from a key Read data from a key
""" """
return self.zk_conn.get(key)[0].decode(self.encoding) path = self.get_schema_path(key)
return self.zk_conn.get(path)[0].decode(self.encoding)
def write(self, kvpairs): def write(self, kvpairs):
""" """
@ -185,26 +215,28 @@ class ZKHandler(object):
key = kvpair[0] key = kvpair[0]
value = kvpair[1] value = kvpair[1]
if not self.exists(key): path = self.get_schema_path(key)
if not self.exists(path):
# Creating a new key # Creating a new key
transaction.create(key, str(value).encode(self.encoding)) transaction.create(path, str(value).encode(self.encoding))
else: else:
# Updating an existing key # Updating an existing key
data = self.zk_conn.get(key) data = self.zk_conn.get(path)
version = data[1].version version = data[1].version
# Validate the expected version after the execution # Validate the expected version after the execution
new_version = version + 1 new_version = version + 1
# Update the data # Update the data
transaction.set_data(key, str(value).encode(self.encoding)) transaction.set_data(path, str(value).encode(self.encoding))
# Check the data # Check the data
try: try:
transaction.check(key, new_version) transaction.check(path, new_version)
except TypeError: except TypeError:
self.log("ZKHandler error: Key '{}' does not match expected version".format(key), state='e') self.log("ZKHandler error: Key '{}' does not match expected version".format(path), state='e')
return False return False
try: try:
@ -222,11 +254,12 @@ class ZKHandler(object):
keys = [keys] keys = [keys]
for key in keys: for key in keys:
if self.exists(key): path = self.get_schema_path(key)
if self.exists(path):
try: try:
self.zk_conn.delete(key, recursive=recursive) self.zk_conn.delete(path, recursive=recursive)
except Exception as e: except Exception as e:
self.log("ZKHandler error: Failed to delete key {}: {}".format(key, e), state='e') self.log("ZKHandler error: Failed to delete key {}: {}".format(path, e), state='e')
return False return False
return True return True
@ -235,7 +268,8 @@ class ZKHandler(object):
""" """
Lists all children of a key Lists all children of a key
""" """
return self.zk_conn.get_children(key) path = self.get_schema_path(key)
return self.zk_conn.get_children(path)
def rename(self, kkpairs): def rename(self, kkpairs):
""" """
@ -247,17 +281,17 @@ class ZKHandler(object):
transaction = self.zk_conn.transaction() transaction = self.zk_conn.transaction()
def rename_element(transaction, source_key, destnation_key): def rename_element(transaction, source_path, destnation_path):
data = self.zk_conn.get(source_key)[0] data = self.zk_conn.get(source_path)[0]
transaction.create(destination_key, data) transaction.create(destination_path, data)
if self.children(source_key): if self.children(source_path):
for child_key in self.children(source_key): for child_path in self.children(source_path):
child_source_key = "{}/{}".format(source_key, child_key) child_source_path = "{}/{}".format(source_path, child_path)
child_destination_key = "{}/{}".format(destination_key, child_key) child_destination_path = "{}/{}".format(destination_path, child_path)
rename_element(transaction, child_source_key, child_destination_key) rename_element(transaction, child_source_path, child_destination_path)
transaction.delete(source_key) transaction.delete(source_path)
for kkpair in (kkpairs): for kkpair in (kkpairs):
if type(kkpair) is not tuple: if type(kkpair) is not tuple:
@ -265,16 +299,19 @@ class ZKHandler(object):
return False return False
source_key = kkpair[0] source_key = kkpair[0]
source_path = self.get_schema_path(source_key)
destination_key = kkpair[1] destination_key = kkpair[1]
destination_path = self.get_schema_path(destination_key)
if not self.exists(source_key): if not self.exists(source_path):
self.log("ZKHander error: Source key '{}' does not exist".format(source_key), state='e') self.log("ZKHander error: Source key '{}' does not exist".format(source_path), state='e')
return False return False
if self.exists(destination_key): if self.exists(destination_path):
self.log("ZKHander error: Destination key '{}' already exists".format(destination_key), state='e') self.log("ZKHander error: Destination key '{}' already exists".format(destination_path), state='e')
return False return False
rename_element(transaction, source_key, destination_key) rename_element(transaction, source_path, destination_path)
try: try:
transaction.commit() transaction.commit()
@ -293,10 +330,12 @@ class ZKHandler(object):
count = 1 count = 1
lock = None lock = None
path = self.get_schema_path(key)
while True: while True:
try: try:
lock_id = str(uuid.uuid1()) lock_id = str(uuid.uuid1())
lock = self.zk_conn.ReadLock(key, lock_id) lock = self.zk_conn.ReadLock(path, lock_id)
break break
except Exception as e: except Exception as e:
if count > 5: if count > 5:
@ -316,10 +355,12 @@ class ZKHandler(object):
count = 1 count = 1
lock = None lock = None
path = self.get_schema_path(key)
while True: while True:
try: try:
lock_id = str(uuid.uuid1()) lock_id = str(uuid.uuid1())
lock = self.zk_conn.WriteLock(key, lock_id) lock = self.zk_conn.WriteLock(path, lock_id)
break break
except Exception as e: except Exception as e:
if count > 5: if count > 5:
@ -339,10 +380,12 @@ class ZKHandler(object):
count = 1 count = 1
lock = None lock = None
path = self.get_schema_path(key)
while True: while True:
try: try:
lock_id = str(uuid.uuid1()) lock_id = str(uuid.uuid1())
lock = self.zk_conn.Lock(key, lock_id) lock = self.zk_conn.Lock(path, lock_id)
break break
except Exception as e: except Exception as e:
if count > 5: if count > 5:
@ -537,8 +580,10 @@ class ZKSchema(object):
return False return False
# Load the schema of a given version from a file # Load the schema of a given version from a file
def load(self, version): def load(self, version, quiet=False):
if not quiet:
print(f'Loading schema version {version}') print(f'Loading schema version {version}')
with open(f'daemon_lib/migrations/versions/{version}.json', 'r') as sfh: with open(f'daemon_lib/migrations/versions/{version}.json', 'r') as sfh:
self.schema = json.load(sfh) self.schema = json.load(sfh)
self.version = self.schema.get('version') self.version = self.schema.get('version')