def _distribute_tables_for_citus(engine): if not getattr(settings, 'ICDS_USE_CITUS', False): return for table, col in DISTRIBUTED_TABLES: with engine.begin() as conn: # TODO: remove this after citus migration res = conn.execute( """ SELECT c.relname AS child FROM pg_inherits JOIN pg_class AS c ON (inhrelid=c.oid) JOIN pg_class as p ON (inhparent=p.oid) where p.relname = %s; """, table ) for child in [row.child for row in res]: # only need this because of reusedb if testing on master and this branch conn.execute('drop table if exists "{}"'.format(child)) create_citus_distributed_table(conn, table, col) for table in REFERENCE_TABLES: with engine.begin() as conn: create_citus_reference_table(conn, table)
def _distribute_citus_tables(apps, schema_editor): with schema_editor.connection.cursor() as cursor: if not is_citus_db(cursor): return for table, col in DISTRIBUTED_TABLES: create_citus_distributed_table(cursor, table, col) for table in REFERENCE_TABLES: create_citus_reference_table(cursor, table)
def _distribute_table(self): config = self.config.sql_settings.citus_config self.session_helper.Session.remove() if not self.session_helper.is_citus_db: # only do this if the database contains the citus extension return from custom.icds_reports.utils.migrations import ( create_citus_distributed_table, create_citus_reference_table ) with self.engine.begin() as connection: if config.distribution_type == 'hash': if config.distribution_column not in self.get_table().columns: raise ColumnNotFoundError("Column '{}' not found.".format(config.distribution_column)) create_citus_distributed_table(connection, self.get_table().name, config.distribution_column) elif config.distribution_type == 'reference': create_citus_reference_table(connection, self.get_table().name) else: raise ValueError("unknown distribution type: %r" % config.distribution_type) return True