async def test_flow():
"""
Test the flow for 3.6 and up, async generators are
not supported in 3.5.
"""
r = RethinkDB()
r.set_loop_type("asyncio")
connection = await r.connect(os.getenv("REBIRTHDB_HOST"))
try:
await r.db_create(INTEGRATION_TEST_DB).run(connection)
except ReqlRuntimeError:
pass
connection.use(INTEGRATION_TEST_DB)
await r.table_create("marvel").run(connection)
marvel_heroes = r.table('marvel')
await marvel_heroes.insert({
'id': 1,
'name': 'Iron Man',
'first_appearance': 'Tales of Suspense #39'
}).run(connection)
cursor = await marvel_heroes.run(connection)
async for hero in cursor:
assert hero['name'] == 'Iron Man'
await connection.close()
def auto(member: Type[Model]):
"""Automatic database and table creation for the given type (Modelchild)."""
if not issubclass(member, Model) or member is Model:
return
rdb = RethinkDB()
conn = rdb.connect(
host=db.HOST,
port=db.PORT,
db=db.DB_NAME,
user=db.USER,
password=db.PASSWORD,
ssl=db.SSL,
timeout=db.TIMEOUT,
)
tables = rdb.table_list().run(conn)
if member.tablename not in tables:
LOG.info("create table %s", member.tablename)
rdb.table_create(member.tablename).run(conn)
indexes = member.get_indexes()
if indexes:
# TODO: at this time, it's only working with simple index
for index in indexes:
rdb.table(member.tablename).index_create(index).run(conn)
rdb.table(member.tablename).index_wait(index).run(conn)
conn.close()
class Iface:
def __init__(self, host="127.0.0.1", port=28015, db="arcos"):
self.r = RethinkDB()
self.conn = self.r.connect(host=host, port=port, db=db)
def clear_table(self, table: str):
self.r.table(table).delete().run(self.conn)
def create_table(self, table: str) -> bool:
try:
self.r.table_create(table).run(self.conn)
return True
except ReqlOpFailedError:
return False
def get_by_uid(self, table: str, uid_name: str, uid_value: str) -> object:
cursor = self.r.table(table).filter(self.r.row[uid_name] == uid_value).run(self.conn)
for document in cursor:
return document
return None
def insert(self, table: str, uid_name: str, uid_value: str, data: object) -> bool:
if self.get_by_uid(table, uid_name, uid_value):
return False
data[uid_name] = uid_value
self.r.table(table).insert(data).run(self.conn)
return True
def update(self, table: str, uid_name: str, uid_value: str, data: object) -> bool:
if not self.get_by_uid(table, uid_name, uid_value):
self.insert(table, uid_name, uid_value, data)
return True
self.r.table(table).filter(self.r.row[uid_name] == uid_value).update(data).run(self.conn)
def connection(self):
r = RethinkDB()
try:
conn = r.connect(host=self.RDB_HOST, port=self.RDB_PORT)
return conn
except RqlDriverError as e:
raise
class BaneSensor:
def __init__(self, ip='localhost', bane_navn='Rink A'):
self.bane_navn = bane_navn
self.r = RethinkDB()
self.conn = self.r.connect(ip, 28015)
def oppdater_sensordata(self, temperatur, luftfuktighet, count=0):
try:
self.r.db(db_navn).table(status_tabell_navn).filter({
'bane':
self.bane_navn
}).update({
'luftfuktighet': luftfuktighet,
'temperatur': temperatur,
}).run(self.conn)
except:
if (count <= 20):
count += 1
self.oppdater_sensordata(temperatur, luftfuktighet, count)
else:
print(
f'{self.bane_navn}: Feilet med å laste opp data til databasen'
)
print(f'{self.bane_navn}: Sensordata oppdatert')
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.prod = True if os.environ.get("pm_id") else False
self.config = cfg
logging.basicConfig(
format="[%(levelname)s] - %(message)s",
level=logging.INFO if os.environ.get("LOG_INFO") else
(logging.INFO if self.prod else logging.DEBUG))
self.logger = logging.getLogger("avabot")
self.session = aiohttp.ClientSession(loop=self.loop)
self.start_time = int(round(time.time() * 1000))
self.uptime = lambda: int(round(time.time() * 1000) - self.start_time)
self.public_dev = False
# We only need to connect to rethink once...
self.r = RethinkDB()
self.r.set_loop_type("asyncio")
self.db_connect_task = self.loop.create_task(self._db_connect())
for cog_name in cog_list:
try:
self.logger.info(f"Loading {cog_name}")
self.load_extension(f"ext.{cog_name}")
except Exception as err:
self.logger.error(f"Failed to load {cog_name}!! {err}")
class Database:
def __init__(self, host, port, user, password, db):
self.host = host
self.port = port
self.user = user
self.password = password
self.db = db
self.r = RethinkDB()
self.conn = None
def query_results(self, query):
"""
queries the database for games matching the given string
"""
results = []
if query == "":
return results
# query for products
self.connect()
cursor = self.r.table('products').filter(
(self.r.row["app_name"].match("(?i)" + query))
& (self.r.row["is_dlc"] == False)).limit(20).run(self.conn)
# build result array
for item in cursor:
if item is not None:
results.append(item)
# close the database connection
self.teardown()
return results
def connect(self):
"""
opens a new connection to the database
"""
try:
# connect to the database
self.conn = self.r.connect(
host=self.host,
port=self.port,
db=self.db,
user=self.user,
password=self.password,
)
except RqlDriverError:
abort(503, "No database connection could be established.")
def teardown(self):
"""
closes an open connection to the database
"""
try:
self.conn.close()
except AttributeError:
pass
def __init__(self):
self.dbname = "DB"
self.r = RethinkDB()
self.connection = self.r.connect("127.0.0.1", "28015").repl()
self.connection.use(self.dbname)
self._init_db()
def __init__(self, host, port, user, password, db):
self.host = host
self.port = port
self.user = user
self.password = password
self.db = db
self.r = RethinkDB()
self.conn = None
def main():
argument_spec = dict(
host=dict(required=True),
port=dict(required=False, type=int, default=28015),
user=dict(required=False, default="admin"),
password=dict(required=False, default="", no_log=True),
ssl=dict(required=False, type=dict, default=None),
table=dict(
required=False,
choices=[
"table_config",
"server_config",
"db_config",
"cluster_config",
"table_status",
"server_status",
"current_issues",
"users",
"permissions",
"jobs",
"stats",
"logs",
],
default="server_status",
),
limit=dict(required=False, type=int, default=10),
)
module = AnsibleModule(
argument_spec=argument_spec,
)
client = RethinkDB()
_params = {
"host": module.params["host"],
"port": module.params["port"],
"user": module.params["user"],
"password": module.params["password"],
"ssl": module.params["ssl"],
"db": "rethinkdb",
}
__res = []
try:
conn = client.connect(**_params)
_res = (
client.table(module.params["table"]).limit(module.params["limit"]).run(conn)
)
while True:
try:
__res.append(_res.next())
except DefaultCursorEmpty:
break
module.exit_json(result=__res)
except (ReqlAuthError, ReqlOpFailedError) as e:
module.fail_json(msg=e.message)
finally:
conn.close(noreply_wait=False)
async def __aenter__(myself):
r = RethinkDB()
r.set_loop_type('asyncio')
if not hasattr(myself, 'conn') or not myself.conn or not myself.conn.is_open():
myself.conn = await r.connect(self.host, self.port, self.db, user=self.user, password=self.password)
self.log.debug(f"Connecting using connection: {myself} (AsyncIO)")
if not myself.conn.is_open():
raise Exception("Cannot open a new rethinkdb connection...")
return myself.conn
def __init__(self, host: str = '127.0.0.1', port: int = 28015) -> None:
self.log = logging
self.format = '%(asctime)s.%(msecs)d | \
%(levelname)s | %(module)s.%(funcName)s:%(lineno)d %(message)s'
self.log.basicConfig(level=logging.DEBUG,
format=self.format,
datefmt='%Y-%m-%d %H:%M:%S')
self.db = RethinkDB()
self._host = host
self._port = port
self.conn = None
def test_database_creation(self):
"""Check, if :meth:rethinkmodel.manage.check_db() creates the database."""
db_name = "test_creation"
config(dbname=db_name)
check_db()
rdb = RethinkDB()
conn = rdb.connect()
dbs = rdb.db_list().run(conn)
if db_name not in dbs:
self.fail("Database named {db_name} was not created")
conn.close()
def _init_db_connection(self):
self.rdb = RethinkDB()
# Create DB connection
try:
self.connection = self.rdb.connect(
host=self.config["data_base"]['host'],
port=self.config["data_base"]['port'],
db=self.config["data_base"]['db_name'])
except RqlDriverError:
self.logger.exception(
"No database connection could be established.")
exit(1)
async def test_tornado_connect(io_loop):
"""
Test the flow for 3.6 and up, async generators are
not supported in 3.5.
"""
r = RethinkDB()
r.set_loop_type("tornado")
connection = await r.connect(os.getenv("REBIRTHDB_HOST"))
dbs = await r.db_list().run(connection)
assert isinstance(dbs, list)
await connection.close()
def __init__(self):
self.dbname = APP_CONFIG['DB']['DB_NAME']
self.tbmovie = 'movies'
self.r = RethinkDB()
try:
self.conn = self.r.connect(host=APP_CONFIG['DB']['DB_HOST'],
db=APP_CONFIG['DB']['DB_NAME'],
user=APP_CONFIG['DB']['DB_USER'],
password=APP_CONFIG['DB']['DB_PASS'],
port=APP_CONFIG['DB']['DB_PORT'])
self.create_db()
except Exception as ex:
print('DB_ERROR', ex)
def __enter__(myself):
r = RethinkDB()
if not hasattr(myself, 'conn') or not myself.conn or not myself.conn.is_open():
myself.conn = r.connect(self.host, self.port, self.db, user=self.user, password=self.password)
self.log.debug(f"Connecting using connection: {myself}")
if not myself.conn.is_open():
raise Exception("Cannot open a new rethinkdb connection...")
self.log.debug(f"Repl-ing connection: {myself}")
myself.conn.repl()
return myself.conn
def test_flow_couroutine_paradigm():
r = RethinkDB()
r.set_loop_type("asyncio")
connection = yield from r.connect(os.getenv("REBIRTHDB_HOST"))
try:
yield from r.db_create(INTEGRATION_TEST_DB).run(connection)
except ReqlRuntimeError:
pass
connection.use(INTEGRATION_TEST_DB)
yield from r.table_create("marvel").run(connection)
marvel_heroes = r.table('marvel')
yield from marvel_heroes.insert({
'id': 1,
'name': 'Iron Man',
'first_appearance': 'Tales of Suspense #39'
}).run(connection)
cursor = yield from marvel_heroes.run(connection)
while (yield from cursor.fetch_next()):
hero = yield from cursor.__anext__()
assert hero['name'] == 'Iron Man'
yield from connection.close()
class PlatformDB(AsyncObject):
'''
Platform database module
'''
async def __init__(self, config={}):
# Fetch configuration from Redis server
host = config.get('host', 'localhost')
port = config.get('port', 28015)
db_name = config.get('db_name', 'platform')
# Connect to RethinkDB
self.db = RethinkDB()
self.db.set_loop_type(library='asyncio')
self.conn = await self.db.connect(host=host,
port=port,
db=db_name)
# Videos table
self.videos = self.db.table('videos')
###########################################################################
# Insert a video
###########################################################################
async def insert_video(self, data):
return await self.videos.insert(data).run(self.conn)
###########################################################################
# Get vidoes - ordered by time
###########################################################################
async def get_videos(self, limit):
return await self.videos.order_by(self.db.desc('timestamp'))\
.limit(limit).run(self.conn)
###########################################################################
# Get count of stored videos
###########################################################################
async def get_videos_count(self):
return await self.videos.count().run(self.conn)
###########################################################################
# Search videos
###########################################################################
async def search_videos(self, query_string, limit):
query_string = '.*%s.*' % query_string.strip().lower()
return await self.videos.filter(
lambda video:
video['title'].downcase().match(query_string) |
video['description'].downcase().match(query_string))\
.limit(limit).order_by(self.db.desc('timestamp')).run(self.conn)
async def __init__(self, config={}):
# Fetch configuration from Redis server
host = config.get('host', 'localhost')
port = config.get('port', 28015)
db_name = config.get('db_name', 'platform')
# Connect to RethinkDB
self.db = RethinkDB()
self.db.set_loop_type(library='asyncio')
self.conn = await self.db.connect(host=host,
port=port,
db=db_name)
# Videos table
self.videos = self.db.table('videos')
def mutate(root, info, id, vms=None, variables=None):
ctx : ContextProtocol = info.context
r = RethinkDB()
table = r.db(opts.database).table(PlaybookLoader.PLAYBOOK_TABLE_NAME)
data = table.get(id).pluck('id').coerce_to('array').run()
if not data:
raise ValueError(f"No such playbook: {id}")
task_id = str(uuid.uuid4())
task_list = PlaybookTaskList()
task_list.upsert_task(ctx.user_authenticator, PlaybookTask(
id=task_id, playbook_id=id, state=PlaybookTaskState.Preparing, message=""))
tornado.ioloop.IOLoop.current().run_in_executor(ctx.executor,
lambda: launch_playbook(ctx, task_id, id, vms, variables))
return PlaybookLaunchMutation(task_id=task_id)
async def db_pool(nursery):
r = RethinkDB()
r.set_loop_type('trio')
db_config = get_config()['database']
db_pool = r.ConnectionPool(host=db_config['host'],
port=db_config['port'],
db='integration_testing',
user=db_config['super_user'],
password=db_config['super_password'],
nursery=nursery)
async with db_pool.connection() as conn:
await r.db_create('integration_testing').run(conn)
yield db_pool
async with db_pool.connection() as conn:
await r.db_drop('integration_testing').run(conn)
await db_pool.close()
def _db_pool(self, nursery):
'''
Create a database connectoin pool.
:param nursery: A Trio nursery to spawn database connections in.
:returns: A RethinkDB connection pool.
'''
r = RethinkDB()
r.set_loop_type('trio')
db_config = self._config['database']
return r.ConnectionPool(host=db_config['host'],
port=db_config['port'],
db=db_config['db'],
user=db_config['user'],
password=db_config['password'],
nursery=nursery)
class Database(object):
def __init__(self, url):
self.r = RethinkDB()
self.conn = self.r.connect(url, 28015)
def client(self):
return self.r
def get_table(self, tbl):
return self.r.db("olympus").table(tbl).run(self.conn)
def get(self, tbl, key):
return self.r.db("olympus").table(tbl).get(key).run(self.conn)
def insert(self, tbl, obj):
self.r.db("olympus").table(tbl).insert(obj, conflict="replace").run(
self.conn)
def connect() -> Tuple[RethinkDB, Any]:
"""Return a RethinkDB object + connection.
You will usually not need to call this function. Rethink:Model use
this function to internally open and close database connection.
"""
rdb = RethinkDB()
connection = rdb.connect(
host=HOST,
port=PORT,
db=DB_NAME,
user=USER,
password=PASSWORD,
timeout=TIMEOUT,
ssl=SSL,
)
return rdb, connection
def __init__(self, database='apscheduler', table='jobs', client=None,
pickle_protocol=pickle.HIGHEST_PROTOCOL, **connect_args):
super(RethinkDBJobStore, self).__init__()
if not database:
raise ValueError('The "database" parameter must not be empty')
if not table:
raise ValueError('The "table" parameter must not be empty')
self.database = database
self.table_name = table
self.table = None
self.client = client
self.pickle_protocol = pickle_protocol
self.connect_args = connect_args
self.r = RethinkDB()
self.conn = None
class DB():
def __init__(self):
self.dbname = APP_CONFIG['DB']['DB_NAME']
self.tbmovie = 'movies'
self.r = RethinkDB()
try:
self.conn = self.r.connect(host=APP_CONFIG['DB']['DB_HOST'],
db=APP_CONFIG['DB']['DB_NAME'],
user=APP_CONFIG['DB']['DB_USER'],
password=APP_CONFIG['DB']['DB_PASS'],
port=APP_CONFIG['DB']['DB_PORT'])
self.create_db()
except Exception as ex:
print('DB_ERROR', ex)
def create_db(self):
if not self.r.db_list().contains(self.dbname).run(self.conn):
self.r.db_create(self.dbname).run(self.conn)
def create_table(self):
if not self.r.db(self.dbname).table_list().contains(self.tbmovie).run(
self.conn):
self.r.db(self.dbname).table_create(self.tbmovie).run(self.conn)
def get_conn(self):
return self.conn
def get_r(self):
return self.r
def execute_query(module):
results = list()
try:
r = RethinkDB()
r.connect(host=module.params.get('host'),
port=module.params.get('port'),
user=module.params.get('user'),
password=module.params.get('password')).repl()
query_result = eval("r.{0}.run()".format(module.params.get('query')))
if type(query_result) is dict:
results.append(query_result)
else:
for document in query_result:
results.append(document)
return results
except Exception as e:
module.fail_json(msg="Error: {0}".format(e))
def main():
argument_spec = dict(
host=dict(required=True),
port=dict(required=False, type=int, default=28015),
user=dict(required=False, default="admin"),
password=dict(required=False, default="", no_log=True),
ssl=dict(required=False, type=dict, default=None),
state=dict(required=False,
choices=["present", "absent"],
default="present"),
database=dict(required=True),
)
module = AnsibleModule(argument_spec=argument_spec, )
client = RethinkDB()
_params = {
"host": module.params["host"],
"port": module.params["port"],
"user": module.params["user"],
"password": module.params["password"],
"ssl": module.params["ssl"],
}
try:
conn = client.connect(**_params)
if module.params["state"].lower() == "present":
_res = client.db_create(module.params["database"]).run(conn)
else:
_res = client.db_drop(module.params["database"]).run(conn)
module.exit_json(changed=True, result=_res)
except ReqlOpFailedError as e:
if module.params["state"].lower(
) == "present" and "already exists" in e.message:
module.exit_json(changed=False, result=e.message)
elif module.params[
"state"] == "absent" and "does not exist" in e.message:
module.exit_json(changed=False, result=e.message)
else:
module.fail_json(msg=e.message)
except ReqlAuthError as e:
module.fail_json(msg=e.message)
finally:
conn.close(noreply_wait=False)
def setup_database():
r = RethinkDB()
r.connect(host='localhost', port=28015).repl()
try:
# Create databases
r.db_create('platform').run()
# Create tables
r.db('platform').table_create('videos', primary_key='video_id').run()
except Exception as e:
print('Database is already setup')
else:
print('Database setup successful')
class Persistence():
def __init__(self):
self.r = RethinkDB()
self.conn = self.r.connect("localhost", 28015)
def insert(self, data):
# add timestamp
mydata = copy.deepcopy(data)
mydata["time"] = int(time.time())
self.last = mydata["time"]
try:
self.r.db('makers-covid').table('global').insert(mydata).run(
self.conn)
except RqlRuntimeError as err:
print(err.message)
def getLast(self):
res = self.r.db('makers-covid').table('global').filter(
self.r.row['time'] >= self.last).run(self.conn)
return res.next()
class RethinkDbDatabaseConnector(DatabaseConnector):
def __init__(self):
self.r = RethinkDB()
self.conn = None
def disconnect(self):
self.conn.close()
def connect(self, hostname, port, database):
self.conn = self.r.connect(hostname, port, database)
return self.r, self.conn
class Rethinkdb:
def __init__(self):
self.__r = RethinkDB()
self.__connector = self.__r.connect(config.host, config.port, db=config.db, user=config.user, password=config.password).repl()
async def closeDB(self):
if self.__connector:
self.__connector.close()
def getR(self):
return self.__r
def launch_playbook(ctx: ContextProtocol, task_id, playbook_id, vms: Optional[List], variables: Optional[Dict[str, Any]]):
with ReDBConnection().get_connection():
r = RethinkDB()
class LaunchPlaybook(Loggable):
"""
This class is used to provide logger for launch_playbook
"""
task_list : PlaybookTaskList
def __init__(self):
self.task_list = PlaybookTaskList()
self.init_log()
def __repr__(self):
return f'PlaybookLauncher <{task_id} ({playbook_id})>'
launcher = LaunchPlaybook()
log = launcher.log
yaml = ruamel.yaml.YAML()
log.debug("Checking access rights for VMs")
def check_access(uuid):
_vm = VM(ctx.user_authenticator, uuid=uuid)
try:
_vm.check_access("playbook")
except XenAdapterUnauthorizedActionException:
launcher.task_list.upsert_task(ctx.user_authenticator, PlaybookTask(
id=task_id, playbook_id=playbook_id, state=PlaybookTaskState.Error,
message=f"VM {_vm}: Access denied (for playbook launcher). Needs 'playbook' access"))
if vms:
for uuid in vms:
check_access(uuid)
else:
vms = []
try:
table = r.db(opts.database).table(PlaybookLoader.PLAYBOOK_TABLE_NAME)
playbook = table.get(playbook_id).run()
temp_dir = tempfile.mkdtemp(prefix='vmemperor-playbook', suffix=playbook_id)
log.debug(f"Creating temporary directory {temp_dir}")
from distutils.dir_util import copy_tree
playbook_dir = playbook['playbook_dir']
log.debug(f"Copying {playbook_dir} into temporary directory")
copy_tree(playbook_dir, temp_dir)
temp_path = Path(temp_dir)
vms_table = r.db(opts.database).table('vms')
documents = vms_table.get_all(*vms).coerce_to('array').run()
if not playbook['inventory']:
hosts_file = 'hosts'
yaml_hosts = {'all': {'hosts': {}}}
for vm in documents:
for interface in vm['interfaces'].values():
network = Network(ctx.user_authenticator, ref=interface['network'])
if network.uuid not in opts.ansible_networks:
log.debug(f"{network} is not a network configured for Ansible. This is probably okay")
continue
if not 'ip' in interface or not interface['ip']:
log.warning(f"Could not get an IP to connect to VM {vm['uuid']}: {network}. This is probably not okay, install Xen drivers")
continue
yaml_hosts['all']['hosts'][vm['name_label']] = {
'ansible_user': '******',
'ansible_host': interface['ip']
}
break
else:
log.warning(
f"Ignoring VM {vm['uuid']}: not connected to any of 'ansible_networks'. Check your configuration")
launcher.task_list.upsert_task(ctx.user_authenticator, PlaybookTask(
id=task_id, playbook_id=playbook_id, state=PlaybookTaskState.ConfigurationWarning,
message=f"Could not connect to VM {vm['uuid']}: Not connected to Ansible network")
)
if yaml_hosts['all']['hosts']:
# Create ansible execution task
with open(temp_path.joinpath(hosts_file), 'w') as file:
yaml.dump(yaml_hosts, file)
log.debug(f"Hosts file created at {file.name}")
else:
log.error(f"No suitable VMs found")
shutil.rmtree(temp_dir, ignore_errors=True)
return
else:
hosts_file = playbook['inventory']
log.debug("Patching variables files...")
for location in playbook['variables_locations']:
this_variables = {}
for variable in playbook['variables_locations'][location]:
value = variables.get(variable, playbook['variables'][variable]['value'])
this_variables[variable] = value
file_name = temp_path.joinpath(location, 'all')
if this_variables:
log.debug(f"Loading file {file_name}")
with open(file_name, 'r') as file:
original_variables = yaml.load(file)
with open(file_name, 'w') as file:
yaml.dump({**original_variables, **this_variables}, file)
log.info(f'File {file_name} patched')
cmd_line = [opts.ansible_playbook, '-i', hosts_file, playbook['playbook']]
cwd = temp_path
log_path = Path(opts.ansible_logs).joinpath(cwd.name)
os.makedirs(log_path)
with open(log_path.joinpath('stdout'), 'w') as _stdout:
with open(log_path.joinpath('stderr'), 'w') as _stderr:
log.debug(f"Running {cmd_line} in {cwd}. Log path: {log_path}")
launcher.task_list.upsert_task(ctx.user_authenticator, PlaybookTask(id=task_id, playbook_id=playbook_id,
state=PlaybookTaskState.Running,
message=f"Task is currently running"))
proc = subprocess.run(cmd_line,
cwd=cwd, stdout=_stdout, stderr=_stderr,
env={"ANSIBLE_HOST_KEY_CHECKING": "False"})
return_code = proc.returncode
launcher.task_list.upsert_task(ctx.user_authenticator,
PlaybookTask(id=task_id, playbook_id=playbook_id,
state=PlaybookTaskState.Finished if return_code == 0 else PlaybookTaskState.Error,
message=f"Task is finished with exit code {return_code}"))
log.info(f'Finished with return code {return_code}. Logs are available in {log_path}')
except Exception as e:
excString = str(e).replace('\n', ' ')
launcher.task_list.upsert_task(ctx.user_authenticator,
PlaybookTask(id=task_id, playbook_id=playbook_id, state=PlaybookTaskState.Error,
message=f"Exception: {excString}"))
log.error(f"Exception: {excString}")
def __init__(self, playbook_name):
from tornado.options import options as opts
self.init_log()
playbook_dir = Path(playbook_name)
playbook_name = playbook_dir.name
self.log.debug(f"Loading playbook {playbook_name} from {playbook_dir}")
from rethinkdb import RethinkDB
r = RethinkDB()
try:
if not playbook_dir.is_dir():
raise ValueError(f'"{playbook_dir.absolute()}" does not exist or not a directory!')
config_file = playbook_dir.joinpath(self._PLAYBOOK_VMEMPEROR_CONF)
if not config_file.is_file():
raise ValueError(f'"{config_file.absolute()}" does not exist or not a file!')
with open(config_file) as file:
config_dict = yaml.safe_load(file)
self.config = config_dict
# Fill optional config parameters
for k,v in self._DEFAULT_CONFIG.items():
if k not in self.config:
self.config[k] = v
# Find variables
keys = self.config['variables'].keys()
self.variables_locations = {}
self.vars = {}
for var in ('host_vars', 'group_vars'):
self.variables_locations[var] = []
host_vars_file = playbook_dir.joinpath(var, 'all')
if not host_vars_file.is_file():
continue
with open(host_vars_file) as file:
host_vars = yaml.safe_load(file)
self.vars[var] = host_vars
for key in keys:
try:
self.config['variables'][key] = {**self.config['variables'][key],
**{'value': host_vars[key]}}
self.variables_locations[var].append(key)
except KeyError:
continue
# Check if playbook file exists
playbook_file = playbook_dir.joinpath(self.config['playbook'])
if not playbook_file.is_file():
raise ValueError(f"Playbook file {playbook_file.absolute()} does not exist")
self.config['playbook_dir'] = str(playbook_dir.absolute())
self.config['id'] = playbook_dir.name
self.config['variables_locations'] = self.variables_locations
table = r.db(opts.database).table(self.PLAYBOOK_TABLE_NAME)
table.insert(self.config).run()
self.log.debug(f"Loaded playbook {self.config['id']}")
except Exception as e:
self.log.error(f"Exception: {e} at {traceback.print_exc()}")
class RethinkDBJobStore(BaseJobStore):
"""
Stores jobs in a RethinkDB database. Any leftover keyword arguments are directly passed to
rethinkdb's `RethinkdbClient <http://www.rethinkdb.com/api/#connect>`_.
Plugin alias: ``rethinkdb``
:param str database: database to store jobs in
:param str collection: collection to store jobs in
:param client: a :class:`rethinkdb.net.Connection` instance to use instead of providing
connection arguments
:param int pickle_protocol: pickle protocol level to use (for serialization), defaults to the
highest available
"""
def __init__(self, database='apscheduler', table='jobs', client=None,
pickle_protocol=pickle.HIGHEST_PROTOCOL, **connect_args):
super(RethinkDBJobStore, self).__init__()
if not database:
raise ValueError('The "database" parameter must not be empty')
if not table:
raise ValueError('The "table" parameter must not be empty')
self.database = database
self.table_name = table
self.table = None
self.client = client
self.pickle_protocol = pickle_protocol
self.connect_args = connect_args
self.r = RethinkDB()
self.conn = None
def start(self, scheduler, alias):
super(RethinkDBJobStore, self).start(scheduler, alias)
if self.client:
self.conn = maybe_ref(self.client)
else:
self.conn = self.r.connect(db=self.database, **self.connect_args)
if self.database not in self.r.db_list().run(self.conn):
self.r.db_create(self.database).run(self.conn)
if self.table_name not in self.r.table_list().run(self.conn):
self.r.table_create(self.table_name).run(self.conn)
if 'next_run_time' not in self.r.table(self.table_name).index_list().run(self.conn):
self.r.table(self.table_name).index_create('next_run_time').run(self.conn)
self.table = self.r.db(self.database).table(self.table_name)
def lookup_job(self, job_id):
results = list(self.table.get_all(job_id).pluck('job_state').run(self.conn))
return self._reconstitute_job(results[0]['job_state']) if results else None
def get_due_jobs(self, now):
return self._get_jobs(self.r.row['next_run_time'] <= datetime_to_utc_timestamp(now))
def get_next_run_time(self):
results = list(
self.table
.filter(self.r.row['next_run_time'] != None) # noqa
.order_by(self.r.asc('next_run_time'))
.map(lambda x: x['next_run_time'])
.limit(1)
.run(self.conn)
)
return utc_timestamp_to_datetime(results[0]) if results else None
def get_all_jobs(self):
jobs = self._get_jobs()
self._fix_paused_jobs_sorting(jobs)
return jobs
def add_job(self, job):
job_dict = {
'id': job.id,
'next_run_time': datetime_to_utc_timestamp(job.next_run_time),
'job_state': self.r.binary(pickle.dumps(job.__getstate__(), self.pickle_protocol))
}
results = self.table.insert(job_dict).run(self.conn)
if results['errors'] > 0:
raise ConflictingIdError(job.id)
def update_job(self, job):
changes = {
'next_run_time': datetime_to_utc_timestamp(job.next_run_time),
'job_state': self.r.binary(pickle.dumps(job.__getstate__(), self.pickle_protocol))
}
results = self.table.get_all(job.id).update(changes).run(self.conn)
skipped = False in map(lambda x: results[x] == 0, results.keys())
if results['skipped'] > 0 or results['errors'] > 0 or not skipped:
raise JobLookupError(job.id)
def remove_job(self, job_id):
results = self.table.get_all(job_id).delete().run(self.conn)
if results['deleted'] + results['skipped'] != 1:
raise JobLookupError(job_id)
def remove_all_jobs(self):
self.table.delete().run(self.conn)
def shutdown(self):
self.conn.close()
def _reconstitute_job(self, job_state):
job_state = pickle.loads(job_state)
job = Job.__new__(Job)
job.__setstate__(job_state)
job._scheduler = self._scheduler
job._jobstore_alias = self._alias
return job
def _get_jobs(self, predicate=None):
jobs = []
failed_job_ids = []
query = (self.table.filter(self.r.row['next_run_time'] != None).filter(predicate) # noqa
if predicate else self.table)
query = query.order_by('next_run_time', 'id').pluck('id', 'job_state')
for document in query.run(self.conn):
try:
jobs.append(self._reconstitute_job(document['job_state']))
except Exception:
self._logger.exception('Unable to restore job "%s" -- removing it', document['id'])
failed_job_ids.append(document['id'])
# Remove all the jobs we failed to restore
if failed_job_ids:
self.r.expr(failed_job_ids).for_each(
lambda job_id: self.table.get_all(job_id).delete()).run(self.conn)
return jobs
def __repr__(self):
connection = self.conn
return '<%s (connection=%s)>' % (self.__class__.__name__, connection)
