def get_db_info():
config = load_config()
db_info = config.get('db_connection_info', {
'user': '******',
'password': '******'
})
return db_info
def safe_get_config():
"""
If working inside of a flask application context,
return the cooked config in the app. Otherwise
load the config from disk and return it (without host
and request context overrides)
"""
if _app_ctx_stack.top:
return current_app.config
else:
log.info("Outside application. Loading config from disk.")
return load_config()
def run_command(args):
cmd = args.cmd
if not cmd:
print "Please enter SQL to run. Example: kitrun.py sqlcmd \"SELECT * FROM tm_players LIMIT 10;\""
return
tier_config = get_tier_config()
service_info = get_service_info()
tier = args.tier or tier_config["tier"]
config = load_config()
tiers = []
if tier == "ALL":
tiers = [t["name"] for t in config["tiers"]]
else:
tiers = [tier]
print "Running SQL Command on Tiers: {}".format(", ".join(tiers))
service_name = service_info["name"]
tenant = args.tenant
tenants = []
for tier_name in tiers:
config = load_config(tier_name)
for t in config.get("tenants", []):
name = t["name"]
if not t.get("db_server"): continue
if tenant and tenant.lower() != name.lower(): continue
t["tier"] = tier_name
tenants.append(t)
for tenant in tenants:
db_server = tenant["db_server"]
tenant_name = tenant["name"]
tier = tenant["tier"]
tenant_name = tenant_name.replace("-{}".format(tier.lower()), "")
full_cmd = "psql postgresql://{db_server}:{port}/{tier}_{tenant}_{service_name} -U postgres -c \"{cmd}\""\
.format(db_server=db_server, tier=tier, tenant=tenant_name, service_name=service_name, cmd=cmd, port=PORT)
print "Running %s" % full_cmd
#! inject the password into env. Highly undesirable
full_cmd = "PGPASSWORD=postgres %s" % full_cmd
os.system(full_cmd)
def tenants_report():
print "The following tenants are registered in config on tier '{}':".format(
get_tier_name())
config = load_config()
for tenant_config in config.get("tenants", []):
name = tenant_config["name"]
# TODO: Get rid of this
if name == "*":
continue
sys.stdout.write(" {}... ".format(name))
db_error = db_check(tenant_config)
if not db_error:
print Fore.GREEN + "OK"
else:
if "does not exist" in db_error:
print Fore.RED + "FAIL! DB does not exist"
else:
print Fore.RED + "Error: %s" % db_error
print "To view more information about each tenant run this command again with the tenant name"
def update_online_statistics():
"""
"""
logger = get_task_logger("update_statistics")
tier_name = get_tier_name()
config = load_config()
tenants = config.get("tenants", [])
logger.info("Updating statistics for %s tenants...", len(tenants))
num_updated = 0
for tenant_config in tenants:
if tenant_config.get("name", "*") == "*":
continue
try:
this_conn_string = get_connection_string(tenant_config, None, tier_name=tier_name)
except TenantNotFoundError:
continue
with sqlalchemy_session(this_conn_string) as session:
result = session.execute("""SELECT COUNT(DISTINCT(player_id)) AS cnt
FROM ck_clients
WHERE heartbeat > NOW() - INTERVAL '1 minutes'""")
cnt = result.fetchone()[0]
if cnt:
num_updated += 1
tenant_name = tenant_config["name"]
name = 'backend.numonline'
row = session.query(Counter).filter(Counter.name == name).first()
if not row:
row = Counter(name=name, counter_type="absolute")
session.add(row)
session.commit()
counter_id = row.counter_id
timestamp = datetime.datetime.utcnow()
add_count(counter_id, 0, timestamp, cnt, is_absolute=True, db_session=session)
session.commit()
print "Updated num_online for %s to %s" % (tenant_name, cnt)
logger.info("Updated %s tenants with online user count", num_updated)
def get_tenants():
tier_name = get_tier_name()
config = load_config()
_tenants = config.get("tenants", [])
tenants = []
for t in _tenants:
t["heartbeat_timeout"] = config.get("heartbeat_timeout",
DEFAULT_HEARTBEAT_TIMEOUT)
if t.get("name", "*") == "*":
continue
try:
this_conn_string = get_connection_string(t,
None,
tier_name=tier_name)
except TenantNotFoundError:
continue
t["conn_string"] = this_conn_string
if config.get("redis_server", None):
t["redis_server"] = config.get("redis_server")
tenants.append(t)
return tenants
def drop_db(tenant, db_host=None, tier_name=None):
config = load_config()
service = config['name']
db_name = construct_db_name(tenant, service, tier_name)
engine = connect(MASTER_DB, db_host)
# disconnect connected clients
engine.execute(
"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{}';"
.format(db_name))
sql = 'DROP DATABASE "{}";'.format(db_name)
engine.execute('COMMIT')
try:
engine.execute(sql)
except Exception as e:
print sql, e
log.info("Tenant %s has been dropped on %s", tenant, db_host
or get_db_info()['server'])
def get_engines():
if conn_string:
engines = {
"dude": {
"engine":
create_engine(conn_string, echo=False,
poolclass=pool.NullPool),
"url":
conn_string
}
}
return engines
engines = {}
from drift.tenant import get_connection_string
config = load_config()
tenants = []
pick_tenant = None
if sys.argv[1] == '-x':
pick_tenant = sys.argv[2]
print 'picking tenant %s' % pick_tenant
for tier in config["tiers"]:
tier_name = tier["name"]
config = load_config(tier_name)
tenant_names = []
for t in config.get("tenants", []):
if not t.get("db_server"):
continue
name = t["name"]
t["tier_name"] = tier_name
if not (pick_tenant and name != pick_tenant) and name != "*":
tenants.append(t)
tenant_names.append(name)
logger.info("Gathering tenants for tier %s: %s", tier_name,
(", ".join(tenant_names) or "(none)"))
db_servers = set([])
for tenant_config in tenants:
from drift.flaskfactory import TenantNotFoundError
try:
conn_info = {"user": "******", "password": "******"}
this_conn_string = get_connection_string(
tenant_config, conn_info, tier_name=tenant_config["tier_name"])
except TenantNotFoundError:
logger.info("Tenant '{}' on tier '{}' not found".format(
tenant_config["name"], tenant_config["tier_name"]))
continue
if this_conn_string not in [e["url"] for e in engines.itervalues()]:
engines["{}.{}".format(tenant_config["tier_name"],
tenant_config["name"])] = rec = {
"url": this_conn_string
}
# quick and dirty connectivity test before trying to upgrade all db's
print "Checking connectivity..."
db_servers = set()
for key, engine in engines.iteritems():
server = engine["url"].split("/")
db_servers.add(server[2].split("@")[1].lower())
err = False
for db_server in db_servers:
port = 5432
sys.stdout.write(db_server + "... ")
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(2)
result = sock.connect_ex((db_server, port))
if result != 0:
print "Unable to connect to server '%s' on port %s" % (db_server,
port)
err = True
else:
print "OK"
if err:
raise Exception(
"Unable to connect to one or more db servers. Bailing out!")
for key in engines.keys():
rec = engines[key]
connection_string = rec["url"]
logger.info("Connecting '{}'...".format(connection_string))
rec['engine'] = create_engine(connection_string,
echo=False,
poolclass=pool.NullPool)
rec['url'] = connection_string
return engines
def run_command(args):
print
from drift import tenant
tenant_name = args.tenant
if not tenant_name:
tenants_report()
return
tier_name = get_tier_name()
config = load_config()
tenant_config = {}
for tenant_config in config.get("tenants", []):
if tenant_config["name"].lower() == tenant_name.lower():
# get the right casing from config
tenant_name = tenant_config["name"]
break
else:
print Fore.RED + "ERROR! Tenant '{}' is not registered in config for tier '{}'" \
.format(tenant_name, tier_name)
print "Please add the tenant into config/config_{}.json and " \
"then run this command again\n".format(tier_name)
return
if not args.action:
tenant_report(tenant_config)
return
db_host = tenant_config["db_server"]
if ":" not in db_host:
db_host += ":{}".format(POSTGRES_PORT)
# TODO validation
db_name = None
if "recreate" in args.action:
actions = ["drop", "create"]
print "Recreating db for tenant '{}'".format(tenant_name)
else:
actions = [args.action]
if "drop" in actions:
print "Dropping tenant {} on {}...".format(tenant_name, db_host)
db_error = db_check(tenant_config)
if db_error:
print "ERROR: You cannot drop the db because it is not reachable: {}".format(
db_error)
return
else:
tenant.drop_db(tenant_name, db_host, tier_name)
if "create" in args.action:
print "Creating tenant '{}' on server '{}'...".format(
tenant_name, db_host)
db_notfound_error = db_check(tenant_config)
if not db_notfound_error:
print "ERROR: You cannot create the database because it already exists"
print "Use the command 'recreate' if you want to drop and create the db"
from drift.tenant import get_connection_string
conn_string = get_connection_string(tenant_config)
print "conn_string = " + conn_string
else:
tenant.create_db(tenant_name, db_host, tier_name)
tenant_report(tenant_config)
def create_db(tenant, db_host=None, tier_name=None):
config = load_config()
service = config['name']
db_name = construct_db_name(tenant, service, tier_name)
username = get_db_info()['user']
engine = connect(MASTER_DB, db_host)
engine.execute('COMMIT')
sql = 'CREATE DATABASE "{}";'.format(db_name)
try:
engine.execute(sql)
except Exception as e:
print sql, e
# TODO: This will only run for the first time and fail in all other cases.
# Maybe test before instead?
sql = 'CREATE ROLE {user} LOGIN PASSWORD "{user}" VALID UNTIL "infinity";'.format(
user=username)
try:
engine.execute(sql)
except Exception as e:
pass
engine = connect(db_name, db_host)
# TODO: Alembic (and sqlalchemy for that matter) don't like schemas. We should
# figure out a way to add these later
models = config.get("models", [])
if not models:
raise Exception("This app has no models defined in config")
for model_module_name in models:
log.info("Building models from %s", model_module_name)
models = importlib.import_module(model_module_name)
models.ModelBase.metadata.create_all(engine)
engine = connect(db_name, db_host)
for schema in schemas:
# Note that this does not automatically grant on tables added later
sql = '''
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA "{schema}" TO {user};
GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA "{schema}" TO {user};
GRANT ALL ON SCHEMA "{schema}" TO {user};'''.format(
schema=schema, user=username)
try:
engine.execute(sql)
except Exception as e:
print sql, e
# stamp the db with the latest alembic upgrade version
ini_path = os.path.join(
os.path.split(os.environ["drift_CONFIG"])[0], "..", "alembic.ini")
alembic_cfg = Config(ini_path)
db_names = alembic_cfg.get_main_option('databases')
connection_string = 'postgresql://%s:%s@%s/%s' % (username, username,
db_host, db_name)
alembic_cfg.set_section_option(db_names, "sqlalchemy.url",
connection_string)
command.stamp(alembic_cfg, "head")
sql = '''
ALTER TABLE alembic_version
OWNER TO postgres;
GRANT ALL ON TABLE alembic_version TO postgres;
GRANT SELECT, UPDATE, INSERT, DELETE ON TABLE alembic_version TO zzp_user;
'''
engine.execute(sql)
return db_name
