def apply_public_func_updates(apps, schema_editor):
path = msh.find_db_functions_dir()
for funcfile in (
# trigger func is the same as original apply (interface not changed) but does not have the drop in the file.
"partitioned_manager_trigger_function.sql", ):
msh.apply_sql_file(schema_editor,
os.path.join(path, funcfile),
literal_placeholder=True)
def reapply_partition_manager_func(apps, schema_editor):
func_file = "partitioned_manager_trigger_function.sql"
func_path = os.path.join(find_db_functions_dir(), func_file)
func_sql = open(func_path).read()
# The file also contains a DROP FUNCTION statement. Let's get rid of that troublesome thing for this migration
func_sql = re.sub(r"DROP FUNCTION IF EXISTS.+",
"-- THE DROP HAS BEEN ERASED --", func_sql)
with schema_editor.connection.cursor() as cur:
cur.execute(func_sql)
def apply_public_function_updates(apps, schema_editor):
path = msh.find_db_functions_dir()
for funcfile in (
"partitioned_tables_manage_trigger_function.sql",
"partitioned_tables_active_trigger_function.sql",
"scan_date_partitions.sql",
):
msh.apply_sql_file(schema_editor,
os.path.join(path, funcfile),
literal_placeholder=True)
def apply_public_func_updates(apps, schema_editor):
path = msh.find_db_functions_dir()
for funcfile in (
# trigger func is the same as original apply (interface not changed) but does not have the drop in the file.
"reapply_partitioned_tables_manage_trigger_function.sql",
# Drop/create here is not being used in triggers as of this migration.
"clone_schema.sql",
):
msh.apply_sql_file(schema_editor,
os.path.join(path, funcfile),
literal_placeholder=True)
def apply_public_function_updates(apps, schema_editor):
path = msh.find_db_functions_dir()
for funcfile in ("jsonb_sha256_text.sql", ):
msh.apply_sql_file(schema_editor, os.path.join(path, funcfile))
class Tenant(TenantMixin):
"""The model used to create a tenant schema."""
# Sometimes the Tenant model can seemingly return funky results,
# so the template schema name is going to get more inline with the
# customer account schema names
_TEMPLATE_SCHEMA = os.environ.get("TEMPLATE_SCHEMA", "template0")
_CLONE_SCHEMA_FUNC_FILENAME = os.path.join(find_db_functions_dir(),
"clone_schema.sql")
_CLONE_SCHEMA_FUNC_SCHEMA = "public"
_CLONE_SHEMA_FUNC_NAME = "clone_schema"
_CLONE_SCHEMA_FUNC_SIG = (
f"{_CLONE_SCHEMA_FUNC_SCHEMA}.{_CLONE_SHEMA_FUNC_NAME}("
"source_schema text, dest_schema text, "
"copy_data boolean DEFAULT false, "
"_verbose boolean DEFAULT false"
")")
# Override the mixin domain url to make it nullable, non-unique
domain_url = None
# Delete all schemas when a tenant is removed
auto_drop_schema = True
def _check_clone_func(self):
LOG.info(
f'Verify that clone function "{self._CLONE_SCHEMA_FUNC_SIG}" exists'
)
res = dbfunc_exists(conn, self._CLONE_SCHEMA_FUNC_SCHEMA,
self._CLONE_SHEMA_FUNC_NAME,
self._CLONE_SCHEMA_FUNC_SIG)
if not res:
LOG.warning(
f'Clone function "{self._CLONE_SCHEMA_FUNC_SIG}" does not exist'
)
LOG.info(
f'Creating clone function "{self._CLONE_SCHEMA_FUNC_SIG}"')
apply_sql_file(conn.schema_editor(),
self._CLONE_SCHEMA_FUNC_FILENAME,
literal_placeholder=True)
res = dbfunc_exists(conn, self._CLONE_SCHEMA_FUNC_SCHEMA,
self._CLONE_SHEMA_FUNC_NAME,
self._CLONE_SCHEMA_FUNC_SIG)
return res
def _verify_template(self, verbosity=1):
LOG.info(
f'Verify that template schema "{self._TEMPLATE_SCHEMA}" exists')
# This is using the teanant table data as the source of truth which can be dangerous.
# If this becomes unreliable, then the database itself should be the source of truth
# and extra code must be written to handle the sync of the table data to the state of
# the database.
template_schema = self.__class__.objects.get_or_create(
schema_name=self._TEMPLATE_SCHEMA)
# Strict check here! Both the record and the schema *should* exist!
return template_schema and schema_exists(self._TEMPLATE_SCHEMA)
def _clone_schema(self):
result = None
with conn.cursor() as cur:
# This db func will clone the schema objects
# bypassing the time it takes to run migrations
sql = """
select public.clone_schema(%s, %s, copy_data => true) as "clone_result";
"""
LOG.info(
f'Cloning template schema "{self._TEMPLATE_SCHEMA}" to "{self.schema_name}" with data'
)
cur.execute(sql, [self._TEMPLATE_SCHEMA, self.schema_name])
result = cur.fetchone()
conn.set_schema_to_public()
return result[0] if result else False
def create_schema(self,
check_if_exists=True,
sync_schema=True,
verbosity=1):
"""
If schema is "public" or matches _TEMPLATE_SCHEMA, then use the superclass' create_schema() method.
Else, verify the template and inputs and use the database clone function.
"""
if self.schema_name in ("public", self._TEMPLATE_SCHEMA):
LOG.info(
f'Using superclass for "{self.schema_name}" schema creation')
return super().create_schema(check_if_exists=True,
sync_schema=sync_schema,
verbosity=verbosity)
# Verify name structure
_check_schema_name(self.schema_name)
# Make sure all of our special pieces are in play
ret = self._check_clone_func()
if not ret:
errmsg = "Missing clone_schema function even after re-applying the function SQL file."
LOG.critical(errmsg)
raise CloneSchemaFuncMissing(errmsg)
ret = self._verify_template(verbosity=verbosity)
if not ret:
errmsg = f'Template schema "{self._TEMPLATE_SCHEMA}" does not exist'
LOG.critical(errmsg)
raise CloneSchemaTemplateMissing(errmsg)
# Always check to see if the schema exists!
if schema_exists(self.schema_name):
LOG.warning(f'Schema "{self.schema_name}" already exists.')
return False
# Clone the schema. The database function will check
# that the source schema exists and the destination schema does not.
self._clone_schema()
LOG.info(
f'Successful clone of "{self._TEMPLATE_SCHEMA}" to "{self.schema_name}"'
)
return True
class Tenant(TenantMixin):
"""The model used to create a tenant schema."""
# Sometimes the Tenant model can seemingly return funky results,
# so the template schema name is going to get more inline with the
# customer account schema names
_TEMPLATE_SCHEMA = os.environ.get("TEMPLATE_SCHEMA", "template0")
_CLONE_SCHEMA_FUNC_FILENAME = os.path.join(find_db_functions_dir(), "clone_schema.sql")
_CLONE_SCHEMA_FUNC_SCHEMA = "public"
_CLONE_SHEMA_FUNC_NAME = "clone_schema"
_CLONE_SCHEMA_FUNC_SIG = (
f"{_CLONE_SCHEMA_FUNC_SCHEMA}.{_CLONE_SHEMA_FUNC_NAME}("
"source_schema text, dest_schema text, "
"copy_data boolean DEFAULT false, "
"_verbose boolean DEFAULT false"
")"
)
# Override the mixin domain url to make it nullable, non-unique
domain_url = None
# Delete all schemas when a tenant is removed
auto_drop_schema = True
auto_create_schema = False
def _check_clone_func(self):
LOG.info(f'Verify that clone function "{self._CLONE_SCHEMA_FUNC_SIG}" exists')
res = dbfunc_exists(
conn, self._CLONE_SCHEMA_FUNC_SCHEMA, self._CLONE_SHEMA_FUNC_NAME, self._CLONE_SCHEMA_FUNC_SIG
)
if not res:
LOG.warning(f'Clone function "{self._CLONE_SCHEMA_FUNC_SIG}" does not exist')
LOG.info(f'Creating clone function "{self._CLONE_SCHEMA_FUNC_SIG}"')
apply_sql_file(conn.schema_editor(), self._CLONE_SCHEMA_FUNC_FILENAME, literal_placeholder=True)
res = dbfunc_exists(
conn, self._CLONE_SCHEMA_FUNC_SCHEMA, self._CLONE_SHEMA_FUNC_NAME, self._CLONE_SCHEMA_FUNC_SIG
)
else:
LOG.info("Clone function exists")
return res
def _verify_template(self, verbosity=1):
LOG.info(f'Verify that template schema "{self._TEMPLATE_SCHEMA}" exists')
# This is using the teanant table data as the source of truth which can be dangerous.
# If this becomes unreliable, then the database itself should be the source of truth
# and extra code must be written to handle the sync of the table data to the state of
# the database.
template_schema, _ = self.__class__.objects.get_or_create(schema_name=self._TEMPLATE_SCHEMA)
try:
template_schema.create_schema()
except Exception as ex:
LOG.error(f"Caught exception {ex.__class__.__name__} during template schema create: {str(ex)}")
raise ex
# Strict check here! Both the record and the schema *should* exist!
res = bool(template_schema) and schema_exists(self._TEMPLATE_SCHEMA)
return res
def _clone_schema(self):
result = None
# This db func will clone the schema objects
# bypassing the time it takes to run migrations
sql = """
select public.clone_schema(%s, %s, copy_data => true) as "clone_result";
"""
LOG.info(f'Cloning template schema "{self._TEMPLATE_SCHEMA}" to "{self.schema_name}"')
LOG.info("Reading catalog for template data")
with conn.cursor() as cur:
cur.execute(sql, [self._TEMPLATE_SCHEMA, self.schema_name])
result = cur.fetchone()
cur.execute("SET search_path = public;")
return result[0] if result else False
# def _clone_schema(self):
# LOG.info("Loading create script from koku_tenant_create.sql file.")
# create_sql_buff = pkgutil.get_data("api.iam", "sql/koku_tenant_create.sql").decode("utf-8")
# LOG.info(f'Cloning template schema "{self._TEMPLATE_SCHEMA}" to "{self.schema_name}"')
# with conn.cursor() as cur:
# cur.execute(f'CREATE SCHEMA IF NOT EXISTS "{self.schema_name}" AUTHORIZATION current_user ;')
# cur.execute(f'SET search_path = "{self.schema_name}", public ;')
# cur.execute(create_sql_buff)
# cur.execute("SET search_path = public ;")
# return True
def create_schema(self, check_if_exists=True, sync_schema=True, verbosity=1):
"""
If schema is "public" or matches _TEMPLATE_SCHEMA, then use the superclass' create_schema() method.
Else, verify the template and inputs and use the database clone function.
"""
if self.schema_name in ("public", self._TEMPLATE_SCHEMA):
LOG.info(f'Using superclass for "{self.schema_name}" schema creation')
return super().create_schema(check_if_exists=True, sync_schema=sync_schema, verbosity=verbosity)
db_exc = None
# Verify name structure
if not _is_valid_schema_name(self.schema_name):
exc = ValidationError(f'Invalid schema name: "{self.schema_name}"')
LOG.error(f"{exc.__class__.__name__}:: {''.join(exc)}")
raise exc
with transaction.atomic():
# Make sure all of our special pieces are in play
ret = self._check_clone_func()
if not ret:
errmsg = "Missing clone_schema function even after re-applying the function SQL file."
LOG.critical(errmsg)
raise CloneSchemaFuncMissing(errmsg)
ret = self._verify_template(verbosity=verbosity)
if not ret:
errmsg = f'Template schema "{self._TEMPLATE_SCHEMA}" does not exist'
LOG.critical(errmsg)
raise CloneSchemaTemplateMissing(errmsg)
# Always check to see if the schema exists!
LOG.info(f"Check if target schema {self.schema_name} already exists")
if schema_exists(self.schema_name):
LOG.warning(f'Schema "{self.schema_name}" already exists. Exit with False.')
return False
# Clone the schema. The database function will check
# that the source schema exists and the destination schema does not.
try:
self._clone_schema()
except Exception as dbe:
db_exc = dbe
LOG.error(
f"""Exception {dbe.__class__.__name__} cloning"""
+ f""" "{self._TEMPLATE_SCHEMA}" to "{self.schema_name}": {str(dbe)}"""
)
LOG.info("Setting transaction to exit with ROLLBACK")
transaction.set_rollback(True) # Set this transaction context to issue a rollback on exit
else:
LOG.info(f'Successful clone of "{self._TEMPLATE_SCHEMA}" to "{self.schema_name}"')
# Set schema to public (even if there was an exception)
with transaction.atomic():
LOG.info("Reset DB search path to public")
conn.set_schema_to_public()
if db_exc:
raise db_exc
return True
def apply_clone_schema(apps, schema_editor):
path = msh.find_db_functions_dir()
msh.apply_sql_file(schema_editor,
os.path.join(path, "clone_schema.sql"),
literal_placeholder=True)
def apply_create_partition_procedure(apps, schema_editor):
path = msh.find_db_functions_dir()
for funcfile in ("create_table_date_range_partition.sql",
"create_date_partitions.sql"):
msh.apply_sql_file(schema_editor, os.path.join(path, funcfile))
def apply_presto_del_trigger_func(apps, schema_editor):
path = msh.find_db_functions_dir()
msh.apply_sql_file(schema_editor,
os.path.join(path, "presto_delete_trigger_func.sql"))
