def partition_registration_list(self, conn, *, num_physical_shards):
"""Method to repartition registration_list for v47 upgrade."""
with conn.cursor() as cursor, utils.db_role_setter(
conn, role_name='dirbs_core_power_user'):
# Create parent partition
cursor.execute("""CREATE TABLE historic_registration_list_new (
LIKE historic_registration_list INCLUDING DEFAULTS
INCLUDING IDENTITY
INCLUDING CONSTRAINTS
INCLUDING STORAGE
INCLUDING COMMENTS
)
PARTITION BY RANGE (virt_imei_shard)
""")
part_utils._grant_perms_registration_list(
conn, part_name='historic_registration_list_new')
# Create child partitions
part_utils.create_imei_shard_partitions(
conn,
tbl_name='historic_registration_list_new',
num_physical_shards=num_physical_shards,
perms_func=part_utils._grant_perms_registration_list,
fillfactor=80)
# Insert data from original partition
cursor.execute("""INSERT INTO historic_registration_list_new
SELECT *
FROM historic_registration_list""")
# Add in indexes to each partition
idx_metadata = [
part_utils.IndexMetadatum(idx_cols=['imei_norm'],
is_unique=True,
partial_sql='WHERE end_date IS NULL')
]
part_utils.add_indices(conn,
tbl_name='historic_registration_list_new',
idx_metadata=idx_metadata)
# Drop old view + table, rename tables, indexes and constraints
cursor.execute('DROP VIEW registration_list')
cursor.execute('DROP TABLE historic_registration_list CASCADE')
part_utils.rename_table_and_indices(
conn,
old_tbl_name='historic_registration_list_new',
new_tbl_name='historic_registration_list',
idx_metadata=idx_metadata)
cursor.execute("""CREATE OR REPLACE VIEW registration_list AS
SELECT imei_norm, make, model, status, virt_imei_shard
FROM historic_registration_list
WHERE end_date IS NULL WITH CHECK OPTION"""
)
cursor.execute("""GRANT SELECT ON registration_list
TO dirbs_core_classify, dirbs_core_api, dirbs_core_import_registration_list"""
)
def _create_index_on_device_id(self, logger, conn):
"""Method to create index on device_id in registration_list."""
logger.info(
'Creating index on device_id in historic_registration_list...')
idx_metadata = [
part_utils.IndexMetadatum(idx_cols=['device_id'],
is_unique=False,
partial_sql='WHERE end_date IS NULL')
]
part_utils.add_indices(conn,
tbl_name='historic_registration_list',
idx_metadata=idx_metadata)
def upgrade(self, conn):
"""Overrides AbstractMigrator upgrade method."""
logger = logging.getLogger('dirbs.db')
logger.info('Creating historic_whitelist table...')
with conn.cursor() as cur:
# create historic table for whitelist
cur.execute(
sql.SQL("""CREATE TABLE historic_whitelist (
imei_norm text NOT NULL,
associated BOOLEAN DEFAULT FALSE,
eir_id text DEFAULT NULL,
start_date TIMESTAMP NOT NULL,
end_date TIMESTAMP DEFAULT NULL,
virt_imei_shard SMALLINT NOT NULL
)
PARTITION BY RANGE (virt_imei_shard)"""))
num_shards = part_utils.num_physical_imei_shards(conn)
logger.debug('Creating Whitelist child partitions...')
part_utils.create_imei_shard_partitions(
conn,
tbl_name='historic_whitelist',
num_physical_shards=num_shards,
fillfactor=80)
# Add indices to each partition
idx_metadata = [
part_utils.IndexMetadatum(idx_cols=['imei_norm'],
is_unique=True,
partial_sql='WHERE end_date IS NULL')
]
part_utils.add_indices(conn,
tbl_name='historic_whitelist',
idx_metadata=idx_metadata)
# creating views to historic_whitelist
cur.execute("""CREATE VIEW whitelist AS
SELECT imei_norm, associated, eir_id, virt_imei_shard
FROM historic_whitelist
WHERE end_date IS NULL WITH CHECK OPTION""")
# create view for imeis that are not associated yet
cur.execute("""CREATE VIEW available_whitelist AS
SELECT imei_norm, virt_imei_shard
FROM historic_whitelist
WHERE associated IS FALSE
AND end_date IS NULL WITH CHECK OPTION""")
# create insert & update trigger on historic_registration_list to update whitelist
# on update and insert
cur.execute(
"""CREATE OR REPLACE FUNCTION insert_whitelist() RETURNS TRIGGER AS
$BODY$
BEGIN
IF new.status = 'whitelist' OR new.status IS NULL THEN
INSERT INTO
historic_whitelist (imei_norm, start_date, end_date, virt_imei_shard)
VALUES (new.imei_norm, new.start_date, new.end_date, new.virt_imei_shard);
END IF;
RETURN new;
END;
$BODY$
LANGUAGE plpgsql;
-- update function
CREATE OR REPLACE FUNCTION update_whitelist() RETURNS TRIGGER AS
$BODY$
BEGIN
UPDATE historic_whitelist
SET end_date = new.end_date
WHERE imei_norm = new.imei_norm
AND new.end_date IS NOT NULL;
RETURN new;
END;
$BODY$
LANGUAGE plpgsql;
-- triggers
CREATE TRIGGER wl_insert_trigger AFTER INSERT ON historic_registration_list
FOR EACH ROW
EXECUTE PROCEDURE insert_whitelist();
CREATE TRIGGER wl_update_trigger AFTER UPDATE ON historic_registration_list
FOR EACH ROW
EXECUTE PROCEDURE update_whitelist();
ALTER TYPE job_command_type RENAME TO job_command_type_old;
--
-- Create type for command
--
CREATE TYPE job_command_type AS ENUM (
'dirbs-catalog',
'dirbs-classify',
'dirbs-db',
'dirbs-import',
'dirbs-listgen',
'dirbs-prune',
'dirbs-report',
'dirbs-whitelist'
);
ALTER TABLE job_metadata ALTER COLUMN command TYPE job_command_type
USING command::TEXT::job_command_type;
DROP TYPE job_command_type_old;
--
-- Whitelist notification triggers
--
CREATE FUNCTION notify_insert_distributor() RETURNS TRIGGER AS
$BODY$
BEGIN
IF new.associated IS FALSE AND new.eir_id IS NULL THEN
PERFORM pg_notify('distributor_updates', row_to_json(NEW)::text);
END IF;
RETURN new;
END;
$BODY$
LANGUAGE plpgsql VOLATILE COST 100;
CREATE FUNCTION notify_remove_distributor() RETURNS TRIGGER AS
$BODY$
BEGIN
IF new.end_date IS NOT NULL THEN
PERFORM pg_notify('distributor_updates', row_to_json(NEW)::text);
END IF;
RETURN new;
END;
$BODY$
LANGUAGE plpgsql VOLATILE COST 100;
CREATE TRIGGER notify_insert_trigger AFTER INSERT ON historic_whitelist
FOR EACH ROW
EXECUTE PROCEDURE notify_insert_distributor();
CREATE TRIGGER notify_remove_trigger AFTER UPDATE ON historic_whitelist
FOR EACH ROW
EXECUTE PROCEDURE notify_remove_distributor();
GRANT SELECT ON historic_whitelist TO dirbs_core_import_registration_list;
GRANT UPDATE ON historic_whitelist TO dirbs_core_import_registration_list;
GRANT INSERT ON historic_whitelist TO dirbs_core_import_registration_list;
GRANT INSERT ON historic_whitelist TO dirbs_core_white_list;
GRANT UPDATE ON historic_whitelist TO dirbs_core_white_list;
GRANT SELECT ON historic_whitelist TO dirbs_core_white_list;
GRANT DELETE ON historic_whitelist TO dirbs_core_white_list;
""") # noqa: Q440, Q449, Q441, Q447
def _repartition_pairing_list(self, conn, *, num_physical_shards):
"""Repartition pairing list to implement change in structure."""
with conn.cursor() as cursor, utils.db_role_setter(
conn, role_name='dirbs_core_power_user'):
cursor.execute("""CREATE TABLE historic_pairing_list_new (
LIKE historic_pairing_list INCLUDING DEFAULTS
INCLUDING IDENTITY
INCLUDING CONSTRAINTS
INCLUDING STORAGE
INCLUDING COMMENTS
)
PARTITION BY RANGE (virt_imei_shard)
""")
cursor.execute(
"""ALTER TABLE historic_pairing_list_new ADD COLUMN msisdn TEXT NOT NULL"""
)
# grant permissions
part_utils._grant_perms_pairing_list(
conn, part_name='historic_pairing_list_new')
# create child partitions
part_utils.create_imei_shard_partitions(
conn,
tbl_name='historic_pairing_list_new',
num_physical_shards=num_physical_shards,
perms_func=part_utils._grant_perms_pairing_list,
fillfactor=80)
# copy data from original partition
cursor.execute("""INSERT INTO historic_pairing_list_new
SELECT p.imei_norm, p.imsi, p.start_date, p.end_date, p.virt_imei_shard, m.msisdn
FROM historic_pairing_list p
INNER JOIN monthly_network_triplets_country m ON p.imsi = m.imsi"""
)
# add indexes
idx_metadata = [
part_utils.IndexMetadatum(
idx_cols=['imei_norm', 'imsi', 'msisdn'],
is_unique=True,
partial_sql='WHERE end_date IS NULL')
]
part_utils.add_indices(conn,
tbl_name='historic_pairing_list_new',
idx_metadata=idx_metadata)
# drop old views, tables, indexes and constraints
cursor.execute('DROP VIEW pairing_list')
cursor.execute('DROP TABLE historic_pairing_list CASCADE')
part_utils.rename_table_and_indices(
conn,
old_tbl_name='historic_pairing_list_new',
new_tbl_name='historic_pairing_list',
idx_metadata=idx_metadata)
# create new view and grant permissions
cursor.execute("""CREATE VIEW pairing_list AS
SELECT imei_norm, imsi, msisdn, virt_imei_shard
FROM historic_pairing_list
WHERE end_date IS NULL WITH CHECK OPTION"""
)
cursor.execute("""GRANT SELECT ON pairing_list
TO dirbs_core_listgen, dirbs_core_report, dirbs_core_api, dirbs_core_import_pairing_list
""")
# drop and recreate staging data insert trigger
cursor.execute("""
DROP FUNCTION pairing_list_staging_data_insert_trigger_fn() CASCADE;
CREATE FUNCTION pairing_list_staging_data_insert_trigger_fn() RETURNS trigger
LANGUAGE plpgsql
AS $$
BEGIN
-- Clean/normalize data before inserting
NEW.imei_norm = normalize_imei(NULLIF(TRIM(NEW.imei), ''));
NEW.imsi = NULLIF(TRIM(new.imsi), '');
NEW.msisdn = NULLIF(TRIM(new.msisdn), '');
RETURN NEW;
END
$$;
""")
def _migrate_device_association_list(self, logger, conn):
"""Method to migrate barred imeis list."""
with conn.cursor() as cursor:
cursor.execute(
sql.SQL("""CREATE TABLE historic_device_association_list (
imei_norm text NOT NULL,
uid text NOT NULL,
start_date TIMESTAMP NOT NULL,
end_date TIMESTAMP,
virt_imei_shard SMALLINT NOT NULL
)
PARTITION BY RANGE (virt_imei_shard);"""))
num_shards = part_utils.num_physical_imei_shards(conn)
logger.debug('Granting permissions to barred_list partitions...')
part_utils._grant_perms_barred_list(
conn, part_name='historic_device_association_list')
logger.debug('Creating barred_list child partitions...')
part_utils.create_imei_shard_partitions(
conn,
tbl_name='historic_device_association_list',
num_physical_shards=num_shards,
perms_func=part_utils._grant_perms_association_list,
fillfactor=80)
# Add indexes to each partition
idx_metadata = [
part_utils.IndexMetadatum(idx_cols=['uid', 'imei_norm'],
is_unique=True,
partial_sql='WHERE end_date IS NULL')
]
part_utils.add_indices(conn,
tbl_name='historic_device_association_list',
idx_metadata=idx_metadata)
# Creating view to historic_barred_list
cursor.execute("""CREATE OR REPLACE VIEW device_association_list AS
SELECT uid, imei_norm, virt_imei_shard
FROM historic_device_association_list
WHERE end_date IS NULL WITH CHECK OPTION"""
) # noqa: Q440
cursor.execute("""GRANT SELECT ON device_association_list
TO dirbs_core_classify, dirbs_core_api,
dirbs_core_import_device_association_list"""
)
# Creating insert trigger function
cursor.execute(
"""CREATE FUNCTION device_association_list_staging_data_insert_trigger_fn() RETURNS TRIGGER
LANGUAGE plpgsql
AS $$
BEGIN
NEW.uid = NULLIF(TRIM(NEW.uid), '');
NEW.imei_norm = normalize_imei(NULLIF(TRIM(NEW.imei), ''));
RETURN NEW;
END
$$;
ALTER FUNCTION device_association_list_staging_data_insert_trigger_fn()
OWNER TO dirbs_core_power_user;
""")
logger.debug(
'Granting create permission to dirbs_core_import_device_association_list...'
)
cursor.execute(
'GRANT CREATE ON SCHEMA core TO dirbs_core_import_device_association_list'
)
def _update_classification_state_job(self, app_config, run_id, curr_date,
virt_imei_range_start,
virt_imei_range_end):
"""Function to update the classificate_state table with IMEIs that are met by this condition (single job)."""
with create_db_connection(app_config.db_config) as conn, conn.cursor(
) as cursor, CodeProfiler() as cp:
src_shard_name = partition_utils.imei_shard_name(
base_name=self.intermediate_tbl_name(run_id),
virt_imei_range_start=virt_imei_range_start,
virt_imei_range_end=virt_imei_range_end)
# Add index on imei_norm
indices = [
partition_utils.IndexMetadatum(idx_cols=['imei_norm'],
is_unique=True)
]
partition_utils.add_indices(conn,
tbl_name=src_shard_name,
idx_metadata=indices)
# Analyze table for better stats/plans
cursor.execute(
sql.SQL('ANALYZE {0}').format(sql.Identifier(src_shard_name)))
# Calculate block date
if curr_date is None:
curr_date = datetime.date.today()
in_amnesty_eval_period, in_amnesty_period = compute_amnesty_flags(
app_config, curr_date)
# If condition is blocking and is not eligible for amnesty, then compute block_date.
# The block_date is set to NULL for amnesty_eligible condition within the eval period.
amnesty_eligible = self.config.amnesty_eligible
sticky_block_date = curr_date + datetime.timedelta(days=self.config.grace_period) \
if self.config.blocking and not (amnesty_eligible and in_amnesty_eval_period) else None
# If the condition's amnesty_eligible flag changed while in eval period, then make sure we update
# the amnesty_granted column in the classification_state table for existing IMEIs meeting that condition.
# These rows will be selected by the existing WHERE clause filters as the block_date would change
# from being NULL to not-NULL or vice-versa.
set_amnesty_granted_column = sql.SQL(', amnesty_granted = {0}').format(sql.Literal(amnesty_eligible)) \
if in_amnesty_eval_period else sql.SQL('')
# If in amnesty period, update the block_date for IMEIs that were previously classified
# as amnesty eligible. This filter is to select amnesty_granted IMEIs with not-NULL block date.
# This is to make sure if the amnesty_end_date was updated, we update the block_date too.
amnesty_block_date_filter = sql.SQL(
'OR cs.amnesty_granted = TRUE'
) if in_amnesty_period else sql.SQL('')
dest_shard_name = partition_utils.imei_shard_name(
base_name='classification_state',
virt_imei_range_start=virt_imei_range_start,
virt_imei_range_end=virt_imei_range_end)
# If a condition is blocking, insert new records into state table with not null blocking date or
# set a not null blocking date for the existing ones having a null block_date.
# Viceversa, if a condition is not blocking, insert new records into state table with Null block_date
# or set a Null block_date for the existing ones having a not-null block_date.
# Set the amnesty_granted column equal to condition's amnesty_eligible flag when in amnesty eval
# period, otherwise always set it to False for new IMEIs meeting the condition.
cursor.execute(
sql.SQL("""INSERT INTO {dest_shard} AS cs(imei_norm,
cond_name,
run_id,
start_date,
end_date,
block_date,
amnesty_granted,
virt_imei_shard)
SELECT imei_norm,
%s,
%s,
%s,
NULL,
%s,
%s,
calc_virt_imei_shard(imei_norm)
FROM {src_shard}
ON CONFLICT (imei_norm, cond_name)
WHERE end_date IS NULL
DO UPDATE
SET block_date = CASE WHEN cs.amnesty_granted = TRUE
AND NOT {in_eval_period}
THEN {amnesty_end_date}
ELSE {sticky_block_date}
END
{set_amnesty_granted_column}
WHERE (cs.block_date IS NULL
AND excluded.block_date IS NOT NULL)
OR (cs.block_date IS NOT NULL
AND excluded.block_date IS NULL)
{amnesty_block_date_filter}"""
) # noqa Q441
.format(src_shard=sql.Identifier(src_shard_name),
dest_shard=sql.Identifier(dest_shard_name),
in_eval_period=sql.Literal(in_amnesty_eval_period),
set_amnesty_granted_column=set_amnesty_granted_column,
amnesty_block_date_filter=amnesty_block_date_filter,
amnesty_end_date=sql.Literal(
app_config.amnesty_config.amnesty_period_end_date),
sticky_block_date=sql.Literal(sticky_block_date)),
[
self.label, run_id, curr_date, sticky_block_date,
(amnesty_eligible and in_amnesty_eval_period)
])
# Get rid of records that no longer exist in the matched IMEIs list
cursor.execute(
sql.SQL("""UPDATE {dest_shard} dst
SET end_date = %s
WHERE cond_name = %s
AND end_date IS NULL
AND NOT EXISTS (SELECT imei_norm
FROM {src_shard}
WHERE imei_norm = dst.imei_norm)"""
).format(src_shard=sql.Identifier(src_shard_name),
dest_shard=sql.Identifier(dest_shard_name)),
[curr_date, self.label])
return cp.duration