def test_get_table_names(self):
test_cases = [
('SELECT * FROM test_a', ['test_a']),
('SELECT * FROM public.test_a', ['test_a']),
('SELECT * FROM "TEST_C"', ['TEST_C']),
('SELECT * FROM public."TEST_C"', ['TEST_C']),
('SELECT * FROM pg_catalog.pg_database', ['pg_database']),
('SELECT rolpassword FROM pg_roles', ['pg_authid']),
('''SELECT p.rolpassword
FROM pg_roles p
JOIN test_b b
ON p.rolpassword = b.id_b''', ['pg_authid', 'test_b']),
('''SELECT id_a, id_b, id_c
FROM (
SELECT *
FROM (
SELECT *
FROM "TEST_C") AS c,
test_b) AS b,
test_a AS a''', ['test_a', 'test_b', 'TEST_C']),
('INSERT INTO test_a VALUES (\'a\')', ['test_a']),
]
context = {
'connection': self.Session.connection()
}
for case in test_cases:
eq_(sorted(datastore_helpers.get_table_names_from_sql(context,
case[0])),
sorted(case[1]))
def search_sql(context, data_dict):
engine = _get_engine(data_dict)
context['connection'] = engine.connect()
timeout = context.get('query_timeout', _TIMEOUT)
_cache_types(context)
sql = data_dict['sql'].replace('%', '%%')
try:
context['connection'].execute(
u'SET LOCAL statement_timeout TO {0}'.format(timeout))
table_names = datastore_helpers.get_table_names_from_sql(context, sql)
log.debug('Tables involved in input SQL: {0!r}'.format(table_names))
system_tables = [t for t in table_names if t.startswith('pg_')]
if len(system_tables):
raise toolkit.NotAuthorized({
'permissions': ['Not authorized to access system tables']
})
results = context['connection'].execute(sql)
return format_results(context, results, data_dict)
except ProgrammingError, e:
if e.orig.pgcode == _PG_ERR_CODE['permission_denied']:
raise toolkit.NotAuthorized({
'permissions': ['Not authorized to read resource.']
})
def _remove_explain(msg):
return (msg.replace('EXPLAIN (FORMAT JSON) ', '')
.replace('EXPLAIN ', ''))
raise ValidationError({
'query': [_remove_explain(str(e))],
'info': {
'statement': [_remove_explain(e.statement)],
'params': [e.params],
'orig': [_remove_explain(str(e.orig))]
}
})
def search_sql(context, data_dict):
engine = _get_engine(data_dict)
context['connection'] = engine.connect()
timeout = context.get('query_timeout', _TIMEOUT)
_cache_types(context)
sql = data_dict['sql'].replace('%', '%%')
try:
context['connection'].execute(
u'SET LOCAL statement_timeout TO {0}'.format(timeout))
table_names = datastore_helpers.get_table_names_from_sql(context, sql)
log.debug('Tables involved in input SQL: {0!r}'.format(table_names))
system_tables = [t for t in table_names if t.startswith('pg_')]
if len(system_tables):
raise toolkit.NotAuthorized({
'permissions': ['Not authorized to access system tables']
})
results = context['connection'].execute(sql)
return format_results(context, results, data_dict)
except ProgrammingError, e:
if e.orig.pgcode == _PG_ERR_CODE['permission_denied']:
raise toolkit.NotAuthorized({
'permissions': ['Not authorized to read resource.']
})
def _remove_explain(msg):
return (msg.replace('EXPLAIN (FORMAT JSON) ', '')
.replace('EXPLAIN ', ''))
raise ValidationError({
'query': [_remove_explain(str(e))],
'info': {
'statement': [_remove_explain(e.statement)],
'params': [e.params],
'orig': [_remove_explain(str(e.orig))]
}
})
def search_sql(context, data_dict):
engine = _get_engine(data_dict)
context["connection"] = engine.connect()
timeout = context.get("query_timeout", _TIMEOUT)
_cache_types(context)
sql = data_dict["sql"].replace("%", "%%")
try:
context["connection"].execute(u"SET LOCAL statement_timeout TO {0}".format(timeout))
table_names = datastore_helpers.get_table_names_from_sql(context, sql)
log.debug("Tables involved in input SQL: {0}".format(table_names))
system_tables = [t for t in table_names if t.startswith("pg_")]
if len(system_tables):
raise toolkit.NotAuthorized({"permissions": ["Not authorized to access system tables"]})
results = context["connection"].execute(sql)
return format_results(context, results, data_dict)
except ProgrammingError, e:
if e.orig.pgcode == _PG_ERR_CODE["permission_denied"]:
raise toolkit.NotAuthorized({"permissions": ["Not authorized to read resource."]})
def _remove_explain(msg):
return msg.replace("EXPLAIN (FORMAT JSON) ", "").replace("EXPLAIN ", "")
raise ValidationError(
{
"query": [_remove_explain(str(e))],
"info": {
"statement": [_remove_explain(e.statement)],
"params": [e.params],
"orig": [_remove_explain(str(e.orig))],
},
}
)
