def _staging_tbl_identifier(self):
"""Staging table name identified."""
return sql.Identifier(self._staging_tbl_name)
def test_join(self):
obj = sql.Composed([sql.Literal("foo"), sql.Identifier("b'ar")])
obj = obj.join(", ")
self.assert_(isinstance(obj, sql.Composed))
self.assertQuotedEqual(obj.as_string(self.conn), "'foo', \"b'ar\"")
def _format_update_survey_resp(american_response, language_field, response):
fmt = _update_survey_response.format(american_response=sql.Literal(american_response),
language_field=sql.Identifier(language_field),
response=sql.Literal(response))
return fmt.as_string(_conn) + '\n'
def test_join(self):
self.assert_(not hasattr(sql.Identifier('foo'), 'join'))
def test_seq(self):
l = [sql.SQL('foo'), sql.Literal('bar'), sql.Identifier('baz')]
self.assertEqual(sql.Composed(l).seq, l)
def _get_contained_items_from_db(self, page, per_page):
cic = self.contained_item_classes[0]
try:
with get_db_cursor() as cursor:
if cic == 'http://linked.data.gov.au/def/gnaf#Address':
id_query = sql.SQL('''SELECT address_detail_pid
FROM {dbschema}.address_detail
ORDER BY address_detail_pid
LIMIT {limit}
OFFSET {offset}''').format(
limit=sql.Literal(per_page),
offset=sql.Literal((page - 1) * per_page),
dbschema=sql.Identifier(config.DB_SCHEMA))
label_prefix = 'Address'
elif cic == 'http://linked.data.gov.au/def/gnaf#Locality':
id_query = sql.SQL('''SELECT locality_pid
FROM {dbschema}.locality
ORDER BY locality_pid
LIMIT {limit}
OFFSET {offset}''').format(
limit=sql.Literal(per_page),
offset=sql.Literal((page - 1) * per_page),
dbschema=sql.Identifier(config.DB_SCHEMA))
label_prefix = 'Locality'
elif cic == 'http://linked.data.gov.au/def/gnaf#StreetLocality':
id_query = sql.SQL('''SELECT street_locality_pid
FROM {dbschema}.street_locality
ORDER BY street_locality_pid
LIMIT {limit}
OFFSET {offset}''').format(
limit=sql.Literal(per_page),
offset=sql.Literal((page - 1) * per_page),
dbschema=sql.Identifier(config.DB_SCHEMA))
label_prefix = 'Street Locality'
elif cic == 'http://linked.data.gov.au/def/gnaf#AddressSite':
id_query = sql.SQL('''SELECT address_site_pid
FROM {dbschema}.address_site
ORDER BY address_site_pid
LIMIT {limit}
OFFSET {offset}''').format(
limit=sql.Literal(per_page),
offset=sql.Literal((page - 1) * per_page),
dbschema=sql.Identifier(config.DB_SCHEMA))
label_prefix = 'Address Site'
else:
raise RuntimeError(
"Cannot get DB objects for unknown contained item class."
)
cursor.execute(id_query)
rows = cursor.fetchall()
for row in rows:
item_pid = row[0]
uri = ''.join([self.uri, item_pid])
label = ' '.join([label_prefix, 'ID:', item_pid])
self.register_items.append((uri, label, item_pid))
except Exception as e:
print(
"Uh oh, can't connect to DB. Invalid dbname, user or password?"
)
print(e)
if args.label_folder is not None:
classdef = {}
with open(os.path.join(args.label_folder, 'classes.txt'), 'r') as f:
lines = f.readlines()
for idx, line in enumerate(lines):
className = line.strip()
# push to database
dbConn.execute(
sql.SQL('''
INSERT INTO {id_lc} (name)
VALUES (
%s
)
ON CONFLICT (name) DO NOTHING;
''').format(id_lc=sql.Identifier(args.project, 'labelclass')),
(className, ))
# get newly assigned index
returnVal = dbConn.execute(
sql.SQL('''
SELECT id FROM {id_lc} WHERE name LIKE %s''').format(
id_lc=sql.Identifier(args.project, 'labelclass')),
(className + '%', ), 1)
classdef[idx] = returnVal[0]['id']
# prepare insertion SQL string
if args.annotation_type == 'annotation':
# get username
usernames = dbConn.execute(
def checkNormal(self):
for table in self.schema:
try:
print(
"\n*****************************************\
\n DB NORMALIZATION CHECKING TABLE:", table,
"\n*****************************************")
# skip DB normalization check if table has no PK
if (len(self.schema[table]['Pk']) == 0):
raise Exception('Skipping table: ', table, ' -> No PK')
# skip DB normalization if any FK references invalid column or table
if (len(self.schema[table]['invalidFks']) > 0):
raise Exception('Skipping table: ', table,
' -> invalid FK: ',
self.schema[table]['invalidFks'])
# if only one NonPrime attribute than table is in 3NF
if (len(self.schema[table]['nonPrimaryAttributes']) == 1):
print('Only one NonPrimeAttr by default in 3NF', )
self.schema[table]['normalized'] = True
continue
# check npa -> npa for 3NF
print('\nALL NPA: ',
self.schema[table]['nonPrimaryAttributes'])
for npa in self.schema[table]['nonPrimaryAttributes']:
# if column has no duplicates than skip it, it passes by default
self.cursor.execute(
sql.SQL("""
SELECT {} FROM {} GROUP BY({}) HAVING count({}) > 1
""").format(sql.Identifier(npa), sql.Identifier(table),
sql.Identifier(npa), sql.Identifier(npa)), )
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
if self.cursor.rowcount == 0:
print("\nNPA: ", npa,
" -> has no duplicates passed 3NF",
"checking rest of table")
self.schema[table]['normalized'] = True
continue
# check all other npa attributes except this one
for checkNpa in self.schema[table]['nonPrimaryAttributes']:
if checkNpa == npa:
continue
print('\n--> checking npa: ', npa)
print('-----> with attr: ', checkNpa)
self.cursor.execute(
sql.SQL("""
DROP TABLE IF EXISTS temp
"""))
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
self.cursor.execute(
sql.SQL("""
DROP TABLE IF EXISTS pairs
"""))
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
self.cursor.execute(
sql.SQL("""
CREATE TABLE IF NOT EXISTS pairs AS (SELECT {} FROM {} GROUP BY {} HAVING COUNT({}) > 1)
""").format(sql.Identifier(npa), sql.Identifier(table),
sql.Identifier(npa), sql.Identifier(npa)), )
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
self.cursor.execute(
sql.SQL("""
CREATE TABLE IF NOT EXISTS temp AS (SELECT {}, {} FROM {} WHERE {} IN (SELECT {} FROM pairs) GROUP BY {}, {})
""").format(sql.Identifier(npa), sql.Identifier(checkNpa),
sql.Identifier(table), sql.Identifier(npa),
sql.Identifier(npa), sql.Identifier(npa),
sql.Identifier(checkNpa)), )
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
self.cursor.execute(
sql.SQL("""
SELECT EXISTS (SELECT {} FROM temp GROUP BY {} HAVING COUNT({}) > 1)
""").format(sql.Identifier(npa), sql.Identifier(npa),
sql.Identifier(npa)), )
finalResult = self.cursor.fetchall()[0][0]
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
if (finalResult == False):
# if not normalized column break out of loop, whole table fails
raise Exception('Column: ', npa, ' -> ', checkNpa,
' failed 3NF')
if (finalResult == True):
print('Column: ', npa, ' -> ', checkNpa,
' passed 3NF checking rest of table')
self.schema[table]['normalized'] = True
except Exception as e:
self.conn.commit()
self.schema[table]['normalized'] = False
print('DB NORMAL COMMANDS FAILED SKIPPING TABLE: ', table,
'\n', e)
continue
def checkRefInt(self):
for table in self.schema:
try:
print(
"\n*********************************\
\n DB REF INT CHECKING TABLE:", table,
"\n*********************************")
# if table has no fk nothing to check
if (len(self.schema[table]['Fks']) == 0):
print('\nDB Ref Int CHECK, Table: ', table,
' -> No fK to check')
self.schema[table]['dbRefInt'] = True
continue
# skip DB check if table has no PK
if (len(self.schema[table]['Pk']) == 0):
raise Exception('Skipping table: ', table, ' -> No PK')
# skip DB if FK references invalid column or table
if (len(self.schema[table]['invalidFks']) > 0):
raise Exception('Skipping table: ', table,
' -> invalid FK: ',
self.schema[table]['invalidFks'])
# check ref Int for remaining tables that have PK and Valid Fk's
for ttc in self.schema[table]['validFks']:
# get table and attr from foreignKey
fKTableName = ttc.split('.')[0]
fKTableAttr = ttc.split('.')[1]
refKey = self.schema[table]['validFks'][ttc]
# get table and attr to check in referenced table
refTableName = self.schema[table]['validFks'][ttc].split(
'.')[0]
refTableAttr = self.schema[table]['validFks'][ttc].split(
'.')[1]
print('-----> checking refInt for fk: ', ttc, ' -> ',
refKey, '\n')
# check if any foreign Key row is NULL
self.cursor.execute(
sql.SQL("""
SELECT * FROM {} WHERE {} IS NULL
""").format(sql.Identifier(fKTableName), sql.Identifier(fKTableAttr)))
totalNullForKeys = self.cursor.rowcount
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
# check if all rows in referencing table
# exist in referenced table
self.cursor.execute(
sql.SQL("""
SELECT * FROM {}
""").format(sql.Identifier(fKTableName)))
totalRows = self.cursor.rowcount
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
self.cursor.execute(
sql.SQL("""
SELECT t.{}, b.{} FROM {} t JOIN {} b ON t.{} = b.{}
""").format(sql.Identifier(fKTableAttr), sql.Identifier(refTableAttr),
sql.Identifier(fKTableName), sql.Identifier(refTableName),
sql.Identifier(fKTableAttr), sql.Identifier(refTableAttr)), )
totalReferencedKeys = self.cursor.rowcount
self.schema[table]['dbCommands'].append(
self.cursor.query.strip())
print('Total Nulls FOR FK in table ', fKTableName, ' : ',
totalNullForKeys)
print('Total Rows in table ', fKTableName, ' : ',
totalRows)
print('Total JOIN Referenced Keys in table ', refTableName,
' : ', totalReferencedKeys)
# failed if we have nulls for foreign keys
if (totalNullForKeys != 0):
raise Exception(
' Failed!!! NULLS : ', totalNullForKeys, ' For:',
fKTableName + '.' + fKTableAttr,
' -> ' + refTableName + "." + refTableAttr)
# if all rows for referencing table
# are not in referenced table, dbRefInt failed
if (totalRows != totalReferencedKeys):
raise Exception(
'Failed!!!', totalRows, ' total rows in table ',
fKTableName, ' Doesnt Match', totalReferencedKeys,
'totalReferencedKeys in', refTableName, ' For:',
fKTableName + '.' + fKTableAttr,
' -> ' + refTableName + "." + refTableAttr)
# if no nulls and row count is equal to referenced keys
if (totalRows == totalReferencedKeys):
print('\nPASSED DB REF INT: ',
fKTableName + '.' + fKTableAttr,
' -> ' + refTableName + "." + refTableAttr)
self.schema[ttc.split('.')[0]]['dbRefInt'] = True
except Exception as e:
# if any DB errors for that table print error
# commit db commands and fail refInt for table
# and skip checking rest of table
self.conn.commit()
print('DB REFINT CHECK FAILED SKIPPING TABLE: ', table, '\n',
e)
self.schema[table]['dbRefInt'] = False
continue
def getDateQueryString(self, project, annotationType, minAge, maxAge,
userNames, skipEmptyImages, goldenQuestionsOnly):
'''
Assembles a DB query string that returns images between a time range.
Useful for reviewing existing annotations.
Inputs:
- minAge: earliest timestamp on which the image(s) have been viewed.
Set to None to leave unrestricted.
- maxAge: latest timestamp of viewing (None = unrestricted).
- userNames: user names to filter the images to. If string, only images
viewed by this respective user are returned. If list, the
images are filtered according to any of the names within.
If None, no user restriction is placed.
- skipEmptyImages: if True, images without an annotation will be ignored.
- goldenQuestionsOnly: if True, images without flag isGoldenQuestion =
True will be ignored.
'''
# column names
fields_anno, _, _ = self._assemble_colnames(annotationType, None)
# user names
usernameString = ''
if userNames is not None:
if isinstance(userNames, str):
usernameString = 'WHERE username = %s'
elif isinstance(userNames, list):
usernameString = 'WHERE username IN %s'
else:
raise Exception('Invalid property for user names')
# date range
timestampString = None
if minAge is not None:
prefix = ('WHERE' if usernameString == '' else 'AND')
timestampString = '{} last_checked::TIMESTAMP > TO_TIMESTAMP(%s)'.format(
prefix)
if maxAge is not None:
prefix = ('WHERE' if (usernameString == ''
and timestampString == '') else 'AND')
timestampString += ' {} last_checked::TIMESTAMP <= TO_TIMESTAMP(%s)'.format(
prefix)
# empty images
if skipEmptyImages:
skipEmptyString = sql.SQL('''
AND image IN (
SELECT image FROM {id_anno}
{usernameString}
)
''').format(id_anno=sql.Identifier(project, 'annotation'),
usernameString=sql.SQL(usernameString))
else:
skipEmptyString = sql.SQL('')
# golden questions
if goldenQuestionsOnly:
goldenQuestionsString = sql.SQL('WHERE isGoldenQuestion = TRUE')
else:
goldenQuestionsString = sql.SQL('')
queryStr = sql.SQL('''
SELECT id, image, cType, username, viewcount, EXTRACT(epoch FROM last_checked) as last_checked, filename, isGoldenQuestion, {annoCols} FROM (
SELECT id AS image, filename, isGoldenQuestion FROM {id_image}
{goldenQuestionsString}
) AS img
JOIN (SELECT image AS iu_image, viewcount, last_checked, username FROM {id_iu}
{usernameString}
{timestampString}
{skipEmptyString}
ORDER BY last_checked ASC
LIMIT %s) AS iu ON img.image = iu.iu_image
LEFT OUTER JOIN (
SELECT id, image AS imID, 'annotation' AS cType, {annoCols} FROM {id_anno} AS anno
{usernameString}
) AS contents ON img.image = contents.imID;
''').format(annoCols=sql.SQL(', ').join(fields_anno),
id_image=sql.Identifier(project, 'image'),
id_iu=sql.Identifier(project, 'image_user'),
id_anno=sql.Identifier(project, 'annotation'),
usernameString=sql.SQL(usernameString),
timestampString=sql.SQL(timestampString),
skipEmptyString=skipEmptyString,
goldenQuestionsString=goldenQuestionsString)
return queryStr
def clean_mail_message(self):
# Clean messages not related to an existing object
_logger.info("Database cleanup for Compassion Switzerland started")
self.env.cr.execute("DELETE FROM mail_message WHERE model IS NULL")
self.env.cr.execute(sql.SQL("SELECT DISTINCT model FROM mail_message"))
models = [m[0] for m in self.env.cr.fetchall()]
for model in models:
table = model.replace(".", "_")
try:
# sql libraries do avoid SQL injection, plus this code
# does not take user input.
# pylint: disable=sql-injection
self.env.cr.execute(
sql.SQL(
"DELETE FROM mail_message WHERE model=%s AND NOT EXISTS ("
" SELECT 1 FROM {} WHERE id = res_id);"
).format(sql.Identifier(table)),
[model],
)
# Commit the cleanup
count = self.env.cr.rowcount
self.env.cr.commit() # pylint: disable=invalid-commit
_logger.info(
"{} mail_message for invalid {} objects removed".format(
count, model
)
)
except ProgrammingError:
self.env.cr.rollback()
# Delete old messages for old objects not useful
today = date.today()
limit = today - relativedelta(years=1)
limit_str = limit.strftime("%Y-%m-%d")
if "account.invoice" in self.env:
self.env.cr.execute(
"DELETE FROM mail_message WHERE model = 'account.invoice' "
"AND res_id IN (SELECT id FROM account_invoice WHERE state IN ("
"'cancel', 'paid') AND date_invoice < %s LIMIT 1000)",
[limit_str],
)
count = self.env.cr.rowcount
self.env.cr.commit() # pylint: disable=invalid-commit
_logger.info("{} old invoice messages removed".format(count))
if "gmc.message.pool" in self.env:
self.env.cr.execute(
"DELETE FROM mail_message WHERE model='gmc.message.pool' AND "
"res_id IN (SELECT id FROM gmc_message_pool WHERE state ="
"'success' LIMIT 1000)"
)
count = self.env.cr.rowcount
self.env.cr.commit() # pylint: disable=invalid-commit
_logger.info("{} old gmc_pool messages removed".format(count))
if "partner.communication.job" in self.env:
self.env.cr.execute(
"DELETE FROM mail_message WHERE model='partner.communication.job' "
"AND res_id IN (SELECT id FROM partner_communication_job WHERE "
"state = 'cancel' OR (state = 'done' AND sent_date < %s)"
"LIMIT 1000)",
[limit_str],
)
count = self.env.cr.rowcount
self.env.cr.commit() # pylint: disable=invalid-commit
_logger.info("{} old partner_communication messages removed".format(count))
if "correspondence" in self.env:
self.env.cr.execute(
"DELETE FROM mail_message WHERE model='correspondence' AND "
"res_id IN (SELECT id FROM correspondence WHERE "
"state IN ('cancel') OR (state = 'done' AND sent_date < %s))",
[limit_str],
)
count = self.env.cr.rowcount
self.env.cr.commit() # pylint: disable=invalid-commit
_logger.info("{} old partner_communication messages removed".format(count))
if "recurring.contract" in self.env:
self.env.cr.execute(
"DELETE FROM mail_message WHERE model='recurring.contract' AND "
"res_id IN (SELECT id FROM recurring_contract WHERE "
"state IN ('terminated', 'cancelled') AND end_date < %s)",
[limit_str],
)
count = self.env.cr.rowcount
self.env.cr.commit() # pylint: disable=invalid-commit
_logger.info("{} old recurring_contract messages removed".format(count))
if "account.analytic.account" in self.env:
self.env.cr.execute(
"DELETE FROM mail_message WHERE model='account.analytic.account' "
"AND date < %s",
[limit_str],
)
count = self.env.cr.rowcount
self.env.cr.commit() # pylint: disable=invalid-commit
_logger.info("{} old analytic_account messages removed".format(count))
def getNextBatchQueryString(self,
project,
annotationType,
predictionType,
order='unlabeled',
subset='default',
demoMode=False):
'''
Assembles a DB query string according to the AL and viewcount ranking criterion.
Inputs:
- order: specifies sorting criterion for request:
- 'unlabeled': prioritize images that have not (yet) been viewed
by the current user (i.e., last_requested timestamp is None or older than 15 minutes
(TODO: make project-specific parameter?))
- 'labeled': put images first in order that have a high user viewcount
- subset: hard constraint on the label status of the images:
- 'default': do not constrain query set
- 'forceLabeled': images must have a viewcount of 1 or more
- 'forceUnlabeled': images must not have been viewed by the current user
- demoMode: set to True to disable sorting criterion and return images in random
order instead.
Note: images market with "isGoldenQuestion" = True will be prioritized if their view-
count by the current user is 0.
'''
# column names
fields_anno, fields_pred, fields_union = self._assemble_colnames(
annotationType, predictionType)
# subset selection fragment
subsetFragment = 'WHERE isGoldenQuestion = FALSE'
subsetFragment_b = ''
orderSpec_a = ''
orderSpec_b = ''
if subset == 'forceLabeled':
subsetFragment = 'WHERE viewcount > 0 AND isGoldenQuestion = FALSE'
subsetFragment_b = 'WHERE viewcount > 0'
elif subset == 'forceUnlabeled':
subsetFragment = 'WHERE (viewcount IS NULL OR viewcount = 0) AND isGoldenQuestion = FALSE'
subsetFragment_b = 'WHERE (viewcount IS NULL OR viewcount = 0)'
if len(subsetFragment):
subsetFragment += ' AND (NOW() - COALESCE(img.last_requested, to_timestamp(0))) > interval \'900 second\''
else:
subsetFragment = 'WHERE (NOW() - COALESCE(img.last_requested, to_timestamp(0))) > interval \'900 second\''
if order == 'unlabeled':
orderSpec_a = 'ORDER BY isgoldenquestion DESC NULLS LAST, viewcount ASC NULLS FIRST, annoCount ASC NULLS FIRST, score DESC NULLS LAST'
orderSpec_b = 'ORDER BY isgoldenquestion DESC NULLS LAST, viewcount ASC NULLS FIRST, annoCount ASC NULLS FIRST, score DESC NULLS LAST'
elif order == 'labeled':
orderSpec_a = 'ORDER BY viewcount DESC NULLS LAST, isgoldenquestion DESC NULLS LAST, score DESC NULLS LAST'
orderSpec_b = 'ORDER BY viewcount DESC NULLS LAST, isgoldenquestion DESC NULLS LAST, score DESC NULLS LAST'
elif order == 'random':
orderSpec_a = 'ORDER BY RANDOM()'
orderSpec_b = 'ORDER BY RANDOM()'
orderSpec_a += ', timeCreated DESC'
orderSpec_b += ', timeCreated DESC'
usernameString = 'WHERE username = %s'
if demoMode:
usernameString = ''
orderSpec_a = ''
orderSpec_b = 'ORDER BY RANDOM()'
gq_user = sql.SQL('')
else:
gq_user = sql.SQL('''AND id NOT IN (
SELECT image FROM {id_iu}
WHERE username = %s
)''').format(id_iu=sql.Identifier(project, 'image_user'))
queryStr = sql.SQL('''
SELECT id, image, cType, viewcount, EXTRACT(epoch FROM last_checked) as last_checked, filename, isGoldenQuestion, {allCols} FROM (
SELECT id AS image, filename, 0 AS viewcount, 0 AS annoCount, NULL AS last_checked, 1E9 AS score, NULL AS timeCreated, isGoldenQuestion FROM {id_img} AS img
WHERE isGoldenQuestion = TRUE
{gq_user}
UNION ALL
SELECT id AS image, filename, viewcount, annoCount, last_checked, score, timeCreated, isGoldenQuestion FROM {id_img} AS img
LEFT OUTER JOIN (
SELECT * FROM {id_iu}
) AS iu ON img.id = iu.image
LEFT OUTER JOIN (
SELECT image, SUM(confidence)/COUNT(confidence) AS score, timeCreated
FROM {id_pred}
WHERE cnnstate = (
SELECT id FROM {id_cnnstate}
ORDER BY timeCreated DESC
LIMIT 1
)
GROUP BY image, timeCreated
) AS img_score ON img.id = img_score.image
LEFT OUTER JOIN (
SELECT image, COUNT(*) AS annoCount
FROM {id_anno}
{usernameString}
GROUP BY image
) AS anno_score ON img.id = anno_score.image
{subset}
{order_a}
LIMIT %s
) AS img_query
LEFT OUTER JOIN (
SELECT id, image AS imID, 'annotation' AS cType, {annoCols} FROM {id_anno} AS anno
{usernameString}
UNION ALL
SELECT id, image AS imID, 'prediction' AS cType, {predCols} FROM {id_pred} AS pred
WHERE cnnstate = (
SELECT id FROM {id_cnnstate}
ORDER BY timeCreated DESC
LIMIT 1
)
) AS contents ON img_query.image = contents.imID
{subset_b}
{order_b};
''').format(id_img=sql.Identifier(project, 'image'),
id_anno=sql.Identifier(project, 'annotation'),
id_pred=sql.Identifier(project, 'prediction'),
id_iu=sql.Identifier(project, 'image_user'),
id_cnnstate=sql.Identifier(project, 'cnnstate'),
gq_user=gq_user,
allCols=sql.SQL(', ').join(fields_union),
annoCols=sql.SQL(', ').join(fields_anno),
predCols=sql.SQL(', ').join(fields_pred),
subset=sql.SQL(subsetFragment),
subset_b=sql.SQL(subsetFragment_b),
order_a=sql.SQL(orderSpec_a),
order_b=sql.SQL(orderSpec_b),
usernameString=sql.SQL(usernameString))
return queryStr
def _postgres_impl(mocked_config):
"""Implementation of fixture to initialise a temporary PostgreSQL instance with a clean DB schema."""
# The system needs to be set to the C locale other than en_US.UTF8 to assume that,
# in collation order uppercase will come before lowercase.
postgresql = testing.postgresql.Postgresql(
initdb_args='-U postgres -A trust --lc-collate=C.UTF-8 '
'--lc-ctype=C.UTF-8')
dsn = postgresql.dsn()
# Monkey-patch Postgres config to use temp postgres instance
for setting in ['database', 'host', 'port', 'user', 'password']:
setattr(mocked_config.db_config, setting, dsn.get(setting, None))
# Run dirbs-db install_roles using db args from the temp postgres instance
runner = CliRunner()
result = runner.invoke(dirbs_db_cli, ['install_roles'],
obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
with create_db_connection(
mocked_config.db_config) as conn, conn.cursor() as cursor:
cursor.execute('CREATE SCHEMA hll;')
cursor.execute('GRANT USAGE ON SCHEMA hll TO dirbs_core_base;')
cursor.execute('CREATE EXTENSION hll SCHEMA hll;')
cursor.execute(
sql.SQL(
'ALTER DATABASE {0} OWNER TO dirbs_core_power_user').format(
sql.Identifier(dsn.get('database'))))
# Run dirbs-db install using db args from the temp postgres instance
result = runner.invoke(dirbs_db_cli, ['install'],
catch_exceptions=False,
obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
# Create the necessary roles for security tests
with create_db_connection(
mocked_config.db_config) as conn, conn.cursor() as cursor:
cursor.execute(
'CREATE ROLE dirbs_import_operator_user IN ROLE dirbs_core_import_operator LOGIN'
)
cursor.execute(
'CREATE ROLE dirbs_import_gsma_user IN ROLE dirbs_core_import_gsma LOGIN'
)
cursor.execute(
'CREATE ROLE dirbs_import_pairing_list_user IN ROLE dirbs_core_import_pairing_list LOGIN'
)
cursor.execute(
'CREATE ROLE dirbs_import_stolen_list_user IN ROLE dirbs_core_import_stolen_list LOGIN'
)
cursor.execute('CREATE ROLE dirbs_import_registration_list_user '
'IN ROLE dirbs_core_import_registration_list LOGIN')
cursor.execute(
'CREATE ROLE dirbs_import_golden_list_user IN ROLE dirbs_core_import_golden_list LOGIN'
)
cursor.execute(
'CREATE ROLE dirbs_import_barred_list_user IN ROLE dirbs_core_import_barred_list LOGIN'
)
cursor.execute(
'CREATE ROLE dirbs_import_barred_tac_list_user IN ROLE dirbs_core_import_barred_tac_list LOGIN'
)
cursor.execute(
'CREATE ROLE dirbs_core_import_subscribers_registration_list_user IN ROLE'
' dirbs_core_import_subscribers_registration_list LOGIN')
cursor.execute(
'CREATE ROLE dirbs_import_device_association_list_user IN ROLE '
'dirbs_core_import_device_association_list LOGIN')
cursor.execute(
'CREATE ROLE dirbs_classify_user IN ROLE dirbs_core_classify LOGIN'
)
cursor.execute(
'CREATE ROLE dirbs_listgen_user IN ROLE dirbs_core_listgen LOGIN')
cursor.execute(
'CREATE ROLE dirbs_report_user IN ROLE dirbs_core_report LOGIN')
cursor.execute(
'CREATE ROLE dirbs_api_user IN ROLE dirbs_core_api LOGIN')
cursor.execute(
'CREATE ROLE dirbs_catalog_user IN ROLE dirbs_core_catalog LOGIN')
cursor.execute(
'CREATE ROLE dirbs_poweruser_login IN ROLE dirbs_core_power_user LOGIN'
)
cursor.execute('CREATE ROLE unknown_user LOGIN')
yield postgresql
postgresql.stop()
def readCSV():
config = StompConfig('tcp://%s:%d' % (host, port), login=user, passcode=password, version='1.1')
client = Stomp(config)
yield client.connect(host='mybroker')
count = 0
start = time.time()
with open(desiredCSV, 'r') as readFile:
csv_reader = csv.reader(readFile)
for row in csv_reader:
if row[4] != 'C' and row[4] != 'G':
try:
cursor.execute(sql.SQL("insert into {} values (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)").format(sql.Identifier('getactivemq')), row)
db_conn.commit()
except:
print "cannot insert into table"
elif row[4] == 'C' or row[4] == 'G':
rowDictionary = {"rowData" : row}
jsonData = json.dumps(rowDictionary)
client.send(destination='atcg', body=jsonData, headers={'persistent': 'false'})
else:
print 'Error reading 5th column'
diff = time.time() - start
print 'Sent %s frames in %f seconds' % (count, diff)
yield client.disconnect(receipt='bye')
def create_schema(cur, schema_name):
print(f'create schema {schema_name}')
schema_id = sql.Identifier(schema_name)
sql_query = sql.SQL('CREATE SCHEMA IF NOT EXISTS {}').format(schema_id)
cur.execute(sql_query)
def write_fields(cur, parent, children, inheritable, weak):
"""
Writes all the inheritable fields of the parent to the children.
Children can be grandchildren etc.
Parameters
----------
cur: psycopg2 cursor
parent: str
The ID of the parent
children: list of str
The IDs of the children
inheritable: list of str
The fields that can be inherited to children
weak: list of str
The fields that can be only be inherited to children if field is empty
Should also be in inheritable
"""
query = sql.SQL("SELECT {} from aen where eventid = %s")
query2 = sql.SQL("UPDATE aen set {} = %s where eventid = %s")
query3 = sql.SQL("UPDATE aen set other = other || %s where eventid = %s")
for col in inheritable:
if col in COLUMNS: # We are working with a name in the columns
cur.execute(query.format((sql.Identifier(col.lower()))),
(parent, ))
value = cur.fetchone()[0]
if value != None: # Make sure we are not removing information
for child in children:
if col in weak: # Only inherit if empty
cur.execute(
query.format((sql.Identifier(col.lower()))),
(child, ))
c_value = cur.fetchone()[0]
if c_value != None: # Something there, so we continue
continue
cur.execute(query2.format((sql.Identifier(col.lower()))), (
value,
child,
))
else: # We are working with a name that could be in other
# Get contents of other
cur.execute(query.format((sql.Identifier('other'))), (parent, ))
other = cur.fetchone()[0]
if other and col in other.keys() and other[col] != '':
value = {col: other[col]}
for child in children:
cur.execute(query.format((sql.Identifier('other'))),
(child, ))
c_other = cur.fetchone()[0]
if c_other and col in c_other.keys(
) and c_other[col] != '':
continue # Something there, so we continue
cur.execute(query3, (
value,
child,
))
def drop_table(cur, table_name):
print(f'drop table {table_name}')
table_id = sql.Identifier(table_name)
sql_query = sql.SQL('DROP TABLE IF EXISTS {}').format(table_id)
cur.execute(sql_query)
def _search_fetch(self, search_detail, search, limit, order):
with_description = 'description' in search_detail['mapping']
results, count = super()._search_fetch(search_detail, search, limit,
order)
if with_description and search:
# Perform search in translations
# TODO Remove when domains will support xml_translate fields
query = sql.SQL("""
SELECT {table}.{id}
FROM {table}
LEFT JOIN ir_ui_view v ON {table}.{view_id} = v.{id}
LEFT JOIN ir_translation t ON v.{id} = t.{res_id}
WHERE t.lang = {lang}
AND t.name = ANY({names})
AND t.type = 'model_terms'
AND t.value ilike {search}
LIMIT {limit}
""").format(
table=sql.Identifier(self._table),
id=sql.Identifier('id'),
view_id=sql.Identifier('view_id'),
res_id=sql.Identifier('res_id'),
lang=sql.Placeholder('lang'),
names=sql.Placeholder('names'),
search=sql.Placeholder('search'),
limit=sql.Placeholder('limit'),
)
self.env.cr.execute(
query, {
'lang': self.env.lang,
'names': ['ir.ui.view,arch_db', 'ir.ui.view,name'],
'search': '%%%s%%' % escape_psql(search),
'limit': limit,
})
ids = {row[0] for row in self.env.cr.fetchall()}
ids.update(results.ids)
domains = search_detail['base_domain'].copy()
domains.append([('id', 'in', list(ids))])
domain = expression.AND(domains)
model = self.sudo() if search_detail.get('requires_sudo') else self
results = model.search(domain,
limit=limit,
order=search_detail.get('order', order))
count = max(count, len(results))
def filter_page(search, page, all_pages):
# Search might have matched words in the xml tags and parameters therefore we make
# sure the terms actually appear inside the text.
text = '%s %s %s' % (page.name, page.url, text_from_html(
page.arch))
pattern = '|'.join(
[re.escape(search_term) for search_term in search.split()])
return re.findall('(%s)' %
pattern, text, flags=re.I) if pattern else False
if 'url' not in order:
results = results._get_most_specific_pages()
if search and with_description:
results = results.filtered(
lambda result: filter_page(search, result, results))
return results, count
try:
conn = psycopg2.connect(dbname='postgres',
user=fediverse_db_user,
host='',
password='')
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
cur = conn.cursor()
print("Creating database " + fediverse_db + ". Please wait...")
cur.execute(
sql.SQL("CREATE DATABASE {}").format(sql.Identifier(fediverse_db)))
print("Database " + fediverse_db + " created!")
except (Exception, psycopg2.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
#############################################################################################
try:
def subscribe(self, channels, condition, callback):
self.callback = callback
# use psycopg2 to get a list of tables in the database
cursor = self.connection.cursor()
cursor.execute("""SELECT table_name FROM information_schema.tables
WHERE table_schema = 'public'""")
tables = [x[0] for x in cursor.fetchall()]
tableDict = {}
#remove uuid and timestamp from the list because those are standard
channel_lower = [x.lower() for x in channels]
if 'uuid' in channel_lower:
channel_lower.remove('uuid')
if 'timestamp' in channel_lower:
channel_lower.remove('timestamp')
wildcard = False
if '*' in channel_lower:
wildcard = True
self.subscribed_channels = channel_lower
# now query each one of those tables to see if the table contains a
# channel that we care about
tableList = []
for table in tables:
tableDict[table] = []
cursor.execute(
"""select column_name from information_schema.columns where table_schema='public' AND table_name=%s""",
(table, ))
columns = [x[0] for x in cursor.fetchall()]
tableDict[table] = columns
channel_set = set(channel_lower)
column_set = set(columns)
intersection = channel_set.intersection(column_set)
if (len(intersection) > 0 or wildcard == True):
tableList.append(table)
tableUUID = {}
uuids = []
for table in tableList:
s = '-'
uuid = s.join(table.split('-')[:-1])
uuids.append(uuid)
if uuid in tableUUID:
tableUUID[uuid].append(table)
else:
tableUUID[uuid] = []
tableUUID[uuid].append(table)
finalTableList = []
if condition != '':
#get column name in condition
later_op = None
split_char = None
if condition.find('=') > -1:
split_char = '='
later_op = operator.eq
elif condition.find('>'):
split_char = '>'
later_op = operator.gt
elif condition.find('<') > -1:
split_char = '<'
later_op = operator.lt
else:
raise ValueError("condition must contain =,>,<")
condition = condition.lower()
condition = "".join(condition.split())
condition_channel = condition.split(split_char)[0]
condition_value = condition.split(split_char)[1]
#look for all tables with the condition we care about
tablesToCheck = []
for key in tableDict:
if condition_channel in tableDict[key]:
tablesToCheck.append(key)
#check if the most recent entry into that table matches the
#condition, if does add it to a good uuid list
goodUUID = []
for table in tablesToCheck:
cursor.execute(
sql.SQL(
"""select {} from {} ORDER BY timestamp DESC limit 1"""
).format(sql.Identifier(condition_channel),
sql.Identifier(table)))
result = [x[0] for x in cursor.fetchall()][0]
if isinstance(result, str):
if later_op(result, condition_value):
#great this is a good uuid
goodUUID.append(s.join(table.split('-')[:-1]))
elif isinstance(result, float):
try:
if later_op(result, float(condition_value)):
#great this is a good uuid
goodUUID.append(s.join(table.split('-')[:-1]))
except:
print('Error retyping condition')
elif isinstance(result, int):
try:
if later_op(result, int(condition_value)):
#great this is a good uuid
goodUUID.append(s.join(table.split('-')[:-1]))
except:
print('Error retyping condition')
#intersect the uuids and good UUID lists to produce a final list
uuid_set = set(uuids)
good_uuid_set = set(goodUUID)
final_set = uuid_set.intersection(good_uuid_set)
for uuid in final_set:
finalTableList = finalTableList + tableUUID[uuid]
else:
finalTableList = tableList
# query all the tables in the final table list
# make a dict with uuid as key and the info the user cares about
# as the values
self.subscribeData = {}
for table in finalTableList:
cursor.execute(
sql.SQL("""select * from {} ORDER BY timestamp DESC limit 1"""
).format(sql.Identifier(table)))
result = cursor.fetchall()[0]
columns = [desc[0] for desc in cursor.description]
if result[columns.index('uuid')] not in self.subscribeData:
self.subscribeData[result[columns.index('uuid')]] = {}
self.subscribeData[result[columns.index('uuid')]]['uuid'] = result[
columns.index('uuid')]
self.subscribeData[result[columns.index(
'uuid')]]['timestamp'] = result[columns.index(
'timestamp')].timestamp() * 1000000
for column in columns:
if (wildcard == True and column != 'timestamp'
and column != 'uuid'):
self.subscribeData[result[columns.index(
'uuid')]][column] = result[columns.index(column)]
elif column in channel_lower:
self.subscribeData[result[columns.index(
'uuid')]][column] = result[columns.index(column)]
#make the initial callbacks to the user with the starting data
for key in self.subscribeData:
self.callback(self.subscribeData[key])
#now turn this list of final tables into a list of MQTT topics
#and subscribe through wave.
s = '-'
for table in finalTableList:
topic = s.join(table.split('-')[:-1]) + '/' + table.split('-')[-1]
self.client.subscribe(self.namespace, topic)
def score_destinations(self,
output_table,
scenario_id=None,
subtract=False,
with_geoms=False,
overwrite=False,
dry=None):
"""
Creates a new db table of scores for each block
args:
output_table -- table to create (optionally schema-qualified)
scenario_id -- the id of the scenario for which scores are calculated
(none means the scores represent the base condition)
subtract -- if true the calculated scores for the scenario represent
a subtraction of that scenario from all other scenarios
overwrite -- overwrite a pre-existing table
dry -- a path to save SQL statements to instead of executing in DB
"""
# make a copy of sql substitutes
subs = dict(self.sql_subs)
# check if a scenarios column exists
if scenario_id is None:
try:
self.get_column_type(self.db_connectivity_table, "scenario")
subs["scenario_where"] = sql.SQL("WHERE scenario IS NULL")
except:
subs["scenario_where"] = sql.SQL("")
else:
subs["scenario_id"] = sql.Literal(scenario_id)
schema, output_table = self.parse_table_name(output_table)
subs["scores_table"] = sql.Identifier(output_table)
if schema is None:
schema = self.get_default_schema()
subs["scores_schema"] = sql.Identifier(schema)
conn = self.get_db_connection()
cur = conn.cursor()
if dry is None:
if overwrite:
self.drop_table(table=output_table, schema=schema, conn=conn)
elif self.table_exists(output_table,
subs["scores_schema"].as_string(conn)):
raise psycopg2.ProgrammingError(
"Table {}.{} already exists".format(
subs["scores_schema"].as_string(conn), output_table))
# create temporary filtered connectivity table
if scenario_id is None:
self._run_sql_script("01_connectivity_table.sql",
subs, ["sql", "destinations"],
conn=conn)
elif subtract:
self._run_sql_script("01_connectivity_table_scenario_subtract.sql",
subs, ["sql", "destinations"],
conn=conn)
else:
self._run_sql_script("01_connectivity_table_scenario.sql",
subs, ["sql", "destinations"],
conn=conn)
# generate high and low stress counts for all categories
print("Counting destinations for each block")
columns = sql.SQL("")
tables = sql.SQL("")
for name, destination in self.destinations.items():
if destination.has_count:
print((" ...{}".format(name)))
destination.count_connections(subs, conn=conn)
destination.calculate_score(subs, conn=conn)
columns += sql.SQL("""
,{table}.hs AS {hs}
,{table}.ls AS {ls}
,{table}.score AS {score}
""").format(
**{
"table": sql.Identifier(destination.workspace_table),
"hs": sql.Identifier(name + "_hs"),
"ls": sql.Identifier(name + "_ls"),
"score": sql.Identifier(name + "_score")
})
tables += sql.SQL("""
LEFT JOIN {schema}.{table} ON blocks.{blocks_id_col} = {table}.block_id
""").format(
**{
"schema": sql.Identifier(destination.workspace_schema),
"table": sql.Identifier(destination.workspace_table),
"blocks_id_col": self.sql_subs["blocks_id_col"]
})
for name, destination in self.destinations.items():
if destination.has_subcats:
columns += sql.SQL("""
,NULL::FLOAT AS {score}
""").format(**{"score": sql.Identifier(name + "_score")})
print("Compiling destination data for all sources into output table")
subs["columns"] = columns
subs["tables"] = tables
self._run_sql_script("04_all_combined.sql",
subs, ["sql", "destinations"],
dry=dry,
conn=conn)
# finally set any category scores
print("Calculating category scores")
self.aggregate_subcategories(self.destinations["overall"],
subs,
conn=conn)
if with_geoms:
self._copy_block_geoms(conn, subs, dry)
conn.commit()
conn.close()
def __init__(self, config, db_connection_string, workspace_schema=None):
"""Sets up a new category of BNA destinations and retrieves data from
the given db table
config : dict
dictionary of config settings (usually from yaml passed to parent BNA object)
db_connection_string : str
string to connect to the database
workspace_schema : str, optional
schema to save interim working tables to
return: None
"""
self.module_dir = os.path.dirname(os.path.abspath(__file__))
DBUtils.__init__(self, db_connection_string)
self.config = config
self.has_subcats = False
self.has_count = False
if "maxpoints" in config:
self.maxpoints = self.config.maxpoints
else:
self.maxpoints = None
if "subcats" in config:
self.has_subcats = True
if "table" in config:
self.has_count = True
if not self.table_exists(config.table):
warnings.warn("No table found for {}".format(config.name))
else:
schema, table = self.parse_table_name(config.table)
if schema is None:
schema = self.get_schema(table)
if "uid" in config:
id_column = config.uid
else:
id_column = self.get_pkid_col(table, schema=schema)
if "geom" in config:
if isinstance(config.geom, str):
geom_col = sql.Identifier(config.geom)
else:
cols = [sql.Identifier(c) for c in config.geom]
geom_col = sql.SQL("COALESCE(") \
+ sql.SQL(",").join([sql.Identifier(c) for c in config.geom]) \
+ sql.SQL(")")
else:
geom_col = sql.SQL("geom")
if "filter" in config:
filter = sql.SQL(config.filter)
else:
filter = sql.Literal(True)
if workspace_schema is None:
self.workspace_schema = "pg_temp"
self.persist = False
else:
self.workspace_schema = workspace_schema
self.persist = True
self.workspace_table = "scores_" + config.name
self.sql_subs = {
"destinations_schema": sql.Identifier(schema),
"destinations_table": sql.Identifier(table),
"destinations_id_col": sql.Identifier(id_column),
"index":
sql.Identifier("idx_{}_blocks".format(config.name)),
"workspace_schema": sql.Identifier(self.workspace_schema),
"workspace_table": sql.Identifier(self.workspace_table),
"destinations_filter": filter
}
if config["method"] == "count":
self.sql_subs["destinations_geom_col"] = geom_col
elif config["method"] == "percentage":
self.sql_subs["val"] = sql.Identifier(config["datafield"])
else:
raise ValueError(
"Unknown scoring method given for {}".format(
config.name))
def test_repr(self):
obj = sql.Composed([sql.Literal("foo"), sql.Identifier("b'ar")])
self.assertEqual(repr(obj),
"""Composed([Literal('foo'), Identifier("b'ar")])""")
self.assertEqual(str(obj), repr(obj))
import traceback
import psycopg2
from psycopg2.extras import execute_values
from psycopg2 import sql
try:
conn = psycopg2.connect("dbname='test' host='127.0.0.1'")
except:
print("I am unable to connect to the database")
cur = conn.cursor()
try:
arr = [("James", 26), ("Kristof", 25)]
execute_values(cur, "INSERT INTO aaa (name, age) VALUES %s", arr)
conn.commit()
tablename = "aaa"
queryText = "INSERT INTO {table} (name, age) VALUES %s"
query = sql.SQL(queryText).format(table=sql.Identifier(tablename))
execute_values(cur, query.as_string(cur), arr)
conn.commit()
except:
print(traceback.format_exc())
def test_eq(self):
l = [sql.Literal("foo"), sql.Identifier("b'ar")]
l2 = [sql.Literal("foo"), sql.Literal("b'ar")]
self.assert_(sql.Composed(l) == sql.Composed(list(l)))
self.assert_(sql.Composed(l) != l)
self.assert_(sql.Composed(l) != sql.Composed(l2))
def tableName(schema, table):
if not schema:
return sql.SQL(""), sql.Identifier(table)
else:
return sql.SQL("{}.").format(sql.Identifier(schema)), sql.Identifier(table)
def test_dict(self):
s = sql.SQL("select {f} from {t}").format(f=sql.Identifier('field'),
t=sql.Identifier('table'))
s1 = s.as_string(self.conn)
self.assert_(isinstance(s1, str))
self.assertEqual(s1, 'select "field" from "table"')
def test_cli_repartition(postgres, mocked_config, db_conn,
operator_data_importer, registration_list_importer,
pairing_list_importer, stolen_list_importer,
gsma_tac_db_importer, tmpdir, logger,
metadata_db_conn, mocked_statsd):
"""Test that the dirbs-db partition script runs without an error."""
import_data(operator_data_importer, 'operator_data', 17, db_conn, logger)
import_data(gsma_tac_db_importer, 'gsma_data', 13, db_conn, logger)
import_data(stolen_list_importer, 'stolen_list', 21, db_conn, logger)
import_data(registration_list_importer, 'registration_list', 20, db_conn,
logger)
import_data(pairing_list_importer, 'pairing_list', 7, db_conn, logger)
# Import second month of operator data to ensure that we have 2 months worth for the same operator
with get_importer(
OperatorDataImporter, db_conn, metadata_db_conn,
mocked_config.db_config, tmpdir, logger, mocked_statsd,
OperatorDataParams(
content='date,imei,imsi,msisdn\n'
'20161201,64220496727231,123456789012345,123456789012345\n'
'20161201,64220496727232,123456789012345,123456789012345',
operator='operator1',
extract=False,
perform_leading_zero_check=False,
perform_unclean_checks=False,
perform_home_network_check=False,
perform_region_checks=False,
perform_historic_checks=False,
mcc_mnc_pairs=[{
'mcc': '111',
'mnc': '04'
}],
perform_file_daterange_check=False)) as new_imp:
expect_success(new_imp, 19, db_conn, logger)
# Run dirbs-classify and dirbs-listgen to populate some tables prior to re-partition
runner = CliRunner()
output_dir = str(tmpdir)
result = runner.invoke(dirbs_classify_cli, ['--no-safety-check'],
obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
result = runner.invoke(dirbs_listgen_cli, [output_dir],
obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
# Make sure that if we run with some invalid phyical shards, that it fails
for num_shards in [-1, 0, 101]:
result = runner.invoke(
dirbs_db_cli,
['repartition', '--num-physical-shards={0:d}'.format(num_shards)],
obj={'APP_CONFIG': mocked_config})
assert result.exit_code != 0
partitioned_tables = [
'classification_state', 'historic_pairing_list',
'historic_registration_list', 'network_imeis',
'monthly_network_triplets_per_mno_operator1_2016_11',
'monthly_network_triplets_country_2016_11', 'blacklist',
'exceptions_lists_operator1', 'notifications_lists_operator1',
'historic_stolen_list'
]
with db_conn, db_conn.cursor() as cursor:
# Manually add one record into the notifications_lists for operator_1 so that the repartitioned table
# is not empty
cursor.execute(
"""INSERT INTO notifications_lists_operator1 (operator_id, imei_norm, imsi, msisdn, block_date,
reasons, start_run_id, end_run_id, delta_reason,
virt_imei_shard)
VALUES ('operator1', '12345678901234', '12345678901234', '1', '20170110',
ARRAY['condition1'], 1125, NULL, 'new', calc_virt_imei_shard('12345678901234'))
""")
# Run dirbs-db repartition to 8 partitions and check that it works
result = runner.invoke(dirbs_db_cli,
['repartition', '--num-physical-shards=8'],
obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
with db_conn, db_conn.cursor() as cursor:
for base_table in partitioned_tables:
cursor.execute(
sql.SQL('SELECT COUNT(*) FROM {0}').format(
sql.Identifier(base_table)))
tbl_count = cursor.fetchone()[0]
assert tbl_count > 0
cursor.execute(
"""SELECT TABLE_NAME
FROM information_schema.tables
WHERE TABLE_NAME LIKE %s
ORDER BY TABLE_NAME""",
['{0}%'.format(base_table)]),
res = [x.table_name for x in cursor]
assert res == [
'{0}'.format(base_table), '{0}_0_12'.format(base_table),
'{0}_13_25'.format(base_table), '{0}_26_38'.format(base_table),
'{0}_39_51'.format(base_table), '{0}_52_63'.format(base_table),
'{0}_64_75'.format(base_table), '{0}_76_87'.format(base_table),
'{0}_88_99'.format(base_table)
]
# Re-partition back to the default 4 shards so that we do not change state for other tests
result = runner.invoke(dirbs_db_cli,
['repartition', '--num-physical-shards=4'],
obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
with db_conn, db_conn.cursor() as cursor:
for base_table in partitioned_tables:
cursor.execute(
sql.SQL('SELECT COUNT(*) FROM {0}').format(
sql.Identifier(base_table)))
tbl_count = cursor.fetchone()[0]
assert tbl_count > 0
cursor.execute(
"""SELECT TABLE_NAME
FROM information_schema.tables
WHERE TABLE_NAME LIKE %s
ORDER BY TABLE_NAME""",
['{0}%'.format(base_table)]),
res = [x.table_name for x in cursor]
assert res == [
'{0}'.format(base_table), '{0}_0_24'.format(base_table),
'{0}_25_49'.format(base_table), '{0}_50_74'.format(base_table),
'{0}_75_99'.format(base_table)
]
def _format_update_question(lang, qid, question):
fmt = _update_question.format(lang=sql.Identifier(lang),
question=sql.Literal(question),
qid=sql.Literal(qid))
return fmt.as_string(_conn) + '\n'
def insert_row(table, cursor, params):
query = sql.SQL('INSERT INTO {} VALUES (%s, %s, %s);').format(
sql.Identifier(table))
query_sql(cursor=cursor, params=params, query=query)