def load_pools(
config: dict,
session: Session = None,
):
pools: dict = config.get("pools", None)
if pools is None:
log.info("No pools found, skipping")
return
log.info("Loading pools from config...")
for key in pools.keys():
val = pools.get(key)
pool = session.query(Pool).filter_by(pool=key).first()
if pool is not None:
log.info(f"Pool exists, skipping: {key}")
continue
log.info("Setting pool: " + key)
pool = Pool(pool=key)
if isinstance(val, dict):
pool.description = val.get("description", "Loaded by zairflow")
pool.slots = val.get("slots", -1)
else:
assert isinstance(val, (int, float))
pool.description = "Loaded from zairflow init"
pool.slots = val or -1
session.add(pool)
session.commit()
def create_pool(name, slots, description, session=None):
"""Create a pool with given parameters."""
if not (name and name.strip()):
raise AirflowBadRequest("Pool name shouldn't be empty")
try:
slots = int(slots)
except ValueError:
raise AirflowBadRequest(f"Bad value for `slots`: {slots}")
# Get the length of the pool column
pool_name_length = Pool.pool.property.columns[0].type.length
if len(name) > pool_name_length:
raise AirflowBadRequest(
f"Pool name can't be more than {pool_name_length} characters")
session.expire_on_commit = False
pool = session.query(Pool).filter_by(pool=name).first()
if pool is None:
pool = Pool(pool=name, slots=slots, description=description)
session.add(pool)
else:
pool.slots = slots
pool.description = description
session.commit()
return pool
def pool(args):
session = settings.Session()
if args.get or (args.set and args.set[0]) or args.delete:
name = args.get or args.delete or args.set[0]
pool = (session.query(Pool).filter(Pool.pool == name).first())
if pool and args.get:
print("{} ".format(pool))
return
elif not pool and (args.get or args.delete):
print("No pool named {} found".format(name))
elif not pool and args.set:
pool = Pool(pool=name, slots=args.set[1], description=args.set[2])
session.add(pool)
session.commit()
print("{} ".format(pool))
elif pool and args.set:
pool.slots = args.set[1]
pool.description = args.set[2]
session.commit()
print("{} ".format(pool))
return
elif pool and args.delete:
session.query(Pool).filter_by(pool=args.delete).delete()
session.commit()
print("Pool {} deleted".format(name))
def pool(args):
session = settings.Session()
if args.get or (args.set and args.set[0]) or args.delete:
name = args.get or args.delete or args.set[0]
pool = (
session.query(Pool)
.filter(Pool.pool == name)
.first())
if pool and args.get:
print("{} ".format(pool))
return
elif not pool and (args.get or args.delete):
print("No pool named {} found".format(name))
elif not pool and args.set:
pool = Pool(
pool=name,
slots=args.set[1],
description=args.set[2])
session.add(pool)
session.commit()
print("{} ".format(pool))
elif pool and args.set:
pool.slots = args.set[1]
pool.description = args.set[2]
session.commit()
print("{} ".format(pool))
return
elif pool and args.delete:
session.query(Pool).filter_by(pool=args.delete).delete()
session.commit()
print("Pool {} deleted".format(name))
def create_pool(name, slots, description, session=None):
"""Create a pool with a given parameters."""
if not (name and name.strip()):
raise PoolBadRequest("Pool name shouldn't be empty")
try:
slots = int(slots)
except ValueError:
raise PoolBadRequest("Bad value for `slots`: %s" % slots)
session.expire_on_commit = False
pool = session.query(Pool).filter_by(pool=name).first()
if pool is None:
pool = Pool(pool=name, slots=slots, description=description)
session.add(pool)
else:
pool.slots = slots
pool.description = description
session.commit()
return pool
def create_pool(name, slots, description, session=None):
"""Create a pool with a given parameters."""
if not (name and name.strip()):
raise AirflowBadRequest("Pool name shouldn't be empty")
try:
slots = int(slots)
except ValueError:
raise AirflowBadRequest("Bad value for `slots`: %s" % slots)
session.expire_on_commit = False
pool = session.query(Pool).filter_by(pool=name).first()
if pool is None:
pool = Pool(pool=name, slots=slots, description=description)
session.add(pool)
else:
pool.slots = slots
pool.description = description
session.commit()
return pool
def generate_init():
"""
Function is checking names of tables in postgres other.
Then information is combined with standard params and pushed as Airflow Variable object.
Also pool is created.
:return:
"""
psql_hook = PostgresHook('airflow_docker_db')
eng = psql_hook.get_sqlalchemy_engine()
df = pd.read_sql(
"select table_name from information_schema.tables where table_schema='other';",
con=eng)
table_list = df['table_name'].tolist()
try:
pool = Pool()
pool.slots = 1
pool.description = 'How many tasks can run at once'
pool.pool = 'generate_tasks'
session = Session()
session.add(pool)
session.commit()
except Exception as ex:
logging.info(f'Could not set pool. Details: {ex}')
init_data = {
'psql_conn_id': 'airflow_docker_db',
'table_list': table_list,
'pool': 'generate_tasks'
}
try:
Variable.set(key='generate_tasks',
value=init_data,
serialize_json=True)
except Exception as ex:
logging.info(f'Could not set global variable. Details: {ex}')
