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"))