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:
parent
76c37e6628
commit
a9a57533a7
|
@ -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):
|
||||||
print(f'Loading schema version {version}')
|
if not quiet:
|
||||||
|
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')
|
||||||
|
|
Loading…
Reference in New Issue