async def _prepare_content_type_tables(self):
"""
Подготавливает соответствие content_type_id и наименование таблицы в БД
"""
logger.info("prepare content type tables")
content_type_table_list = await self._dst_database.fetch_raw_sql(
SQLRepository.get_content_type_table_sql()
)
content_type_table_dict = {
(app_label, model): table_name
for table_name, app_label, model in content_type_table_list
}
content_type_list = await self._src_database.fetch_raw_sql(
SQLRepository.get_content_type_sql()
)
content_type_dict = {
(app_label, model): content_type_id
for content_type_id, app_label, model in content_type_list
}
for key in content_type_table_dict.keys():
self.content_type_table[content_type_table_dict[key]] = (
content_type_dict[key]
)
del content_type_table_list[:]
del content_type_table_dict
del content_type_list[:]
del content_type_dict
async def set_max_sequence(self, dst_pool: Pool):
async with dst_pool.acquire() as connection:
try:
get_serial_sequence_sql = SQLRepository.get_serial_sequence_sql(
table_name=self.name,
pk_column_name=self.primary_key.name,
)
except AttributeError:
logger.warning(
f'AttributeError --- {self.name} --- set_max_sequence')
return
serial_seq_name = await connection.fetchrow(get_serial_sequence_sql
)
if serial_seq_name and serial_seq_name[0]:
serial_seq_name = serial_seq_name[0]
max_val = self.max_pk + 100000
set_sequence_val_sql = (
SQLRepository.get_set_sequence_value_sql(
seq_name=serial_seq_name,
seq_val=max_val,
))
await connection.execute(set_sequence_val_sql)
async def enable(self):
create_fdw_extension_sql = SQLRepository.get_create_fdw_extension_sql()
create_server_sql = SQLRepository.get_create_server_sql(
src_host=self._src_database.db_connection_parameters.host,
src_port=self._src_database.db_connection_parameters.port,
src_dbname=self._src_database.db_connection_parameters.dbname,
)
create_user_mapping_sql = SQLRepository.get_create_user_mapping_sql(
dst_user=self._dst_database.db_connection_parameters.user,
src_user=self._src_database.db_connection_parameters.user,
src_password=self._src_database.db_connection_parameters.password,
)
create_temp_src_schema_sql = (
SQLRepository.get_create_temp_src_schema_sql(
dst_user=self._dst_database.db_connection_parameters.user,
)
)
import_foreign_schema_sql = SQLRepository.get_import_foreign_schema_sql(
src_schema=self._src_database.db_connection_parameters.schema,
tables=self._dst_database.table_names,
)
async with self._dst_pool.acquire() as connection:
await asyncio.wait([connection.execute(create_fdw_extension_sql)])
await asyncio.wait([connection.execute(create_server_sql)])
await asyncio.wait([connection.execute(create_user_mapping_sql)])
await asyncio.wait([connection.execute(create_temp_src_schema_sql)])
await asyncio.wait([connection.execute(import_foreign_schema_sql)])
async def disable(self):
drop_temp_src_schema_sql = SQLRepository.get_drop_temp_src_schema_sql()
drop_user_mapping_sql = SQLRepository.get_drop_user_mapping_sql(
dst_user=self._dst_database.db_connection_parameters.user,
)
drop_fdw_extension_sql = SQLRepository.get_drop_fdw_extension_sql()
async with self._dst_pool.acquire() as connection:
await asyncio.wait([connection.execute(drop_temp_src_schema_sql)])
await asyncio.wait([connection.execute(drop_user_mapping_sql)])
await asyncio.wait([connection.execute(drop_fdw_extension_sql)])
async def truncate_tables(self):
"""
Truncating tables
"""
if settings.IS_TRUNCATE_TABLES:
logger.info('start truncating tables..')
if settings.TABLES_TRUNCATE_INCLUDED:
table_names = settings.TABLES_TRUNCATE_INCLUDED
else:
table_names = tuple(
filter(
lambda table_name: (table_name not in settings.
TABLES_WITH_GENERIC_FOREIGN_KEY),
self.table_names,
))
if settings.TABLES_TRUNCATE_EXCLUDED:
table_names = tuple(
filter(
lambda table_name:
(table_name not in settings.TABLES_TRUNCATE_EXCLUDED),
table_names,
))
truncate_table_queries = SQLRepository.get_truncate_table_queries(
table_names=table_names, )
for query in truncate_table_queries:
await self.execute_raw_sql(query)
logger.info('truncating tables finished.')
async def _transfer_chunk_table_data(
self,
table: DBTable,
need_import_ids_chunk: List[Union[int, str]],
):
"""
Порционный перенос данных таблицы в целевую БД
"""
transfer_sql = SQLRepository.get_transfer_records_sql(
table=table,
connection_params_str=self._src_database.connection_str,
primary_key_ids=need_import_ids_chunk,
)
logger.info(f'transfer chunk table data - {table.name}')
transferred_ids = None
async with self._dst_database.connection_pool.acquire() as connection:
try:
transferred_ids = await connection.fetch(transfer_sql)
except (
UndefinedColumnError,
NotNullViolationError,
PostgresSyntaxError,
) as e:
raise PostgresError(
f'{str(e)}, table - {table.name}, '
f'sql - {transfer_sql} --- _transfer_chunk_table_data')
if transferred_ids:
transferred_ids = [tr[0] for tr in transferred_ids]
table.transferred_pks.update(transferred_ids)
del transfer_sql
async def _set_table_counters(self, table_name: str):
"""
Filling table max pk and count of records
"""
async with self._src_database.connection_pool.acquire() as connection:
table = self._dst_database.tables[table_name]
try:
count_table_records_sql = (
SQLRepository.get_count_table_records(
primary_key=table.primary_key, ))
except AttributeError as e:
logger.warning(f'{str(e)} --- _set_table_counters {"-"*10} - '
f"{table.name}")
raise AttributeError
except UndefinedFunctionError:
raise UndefinedFunctionError
res = await connection.fetchrow(count_table_records_sql)
if res and res[0] and res[1]:
logger.debug(f"table {table_name} with full count {res[0]}, "
f"max pk - {res[1]}")
table.full_count = int(res[0])
table.max_pk = (int(res[1]) if isinstance(res[1], int) else
table.full_count + 100000)
del count_table_records_sql
async def enable_triggers(self):
"""
Enable database triggers
"""
enable_triggers_sql = SQLRepository.get_enable_triggers_sql()
await self.execute_raw_sql(enable_triggers_sql)
logger.info('triggers enabled.')
async def disable_triggers(self):
"""
Disable database triggers
"""
disable_triggers_sql = SQLRepository.get_disable_triggers_sql()
await self.execute_raw_sql(disable_triggers_sql)
logger.info('trigger disabled.')
async def prepare_table_names(self):
"""
Preparing database table names list
"""
select_tables_names_list_sql = SQLRepository.get_select_tables_names_list_sql( # noqa
excluded_tables=settings.EXCLUDED_TABLES, )
async with self._connection_pool.acquire() as connection:
table_names = await connection.fetch(
query=select_tables_names_list_sql, )
self.table_names = [
table_name_rec[0] for table_name_rec in table_names
]
async def _prepare_chunk_tables(
self,
chunk_table_names: Iterable[str],
):
"""
Preparing tables of chunk table names
"""
getting_tables_columns_sql = SQLRepository.get_table_columns_sql(
table_names=make_str_from_iterable(
iterable=chunk_table_names,
with_quotes=True,
quote='\'',
), )
async with self._connection_pool.acquire() as connection:
records = await connection.fetch(
query=getting_tables_columns_sql, )
coroutines = [
self.tables[table_name].append_column(
column_name=column_name,
data_type=data_type,
ordinal_position=ordinal_position,
constraint_table=self.tables.get(constraint_table_name),
constraint_type=constraint_type,
) for (
table_name,
column_name,
data_type,
ordinal_position,
constraint_table_name,
constraint_type,
) in records
]
if coroutines:
await asyncio.gather(*coroutines)
self.clear_cache()