def query(self,
query,
data=None,
fmt='PrettyCompact',
stream=False,
verbose=False,
query_id=None,
compress=False,
**kwargs):
if query.lstrip()[:6].upper().startswith('INSERT'):
query_split = query.split()
else:
query = sqlparse.format(query, strip_comments=True).rstrip(';')
if verbose and self.cli_settings.get('show_formatted_query'):
# Highlight & reformat the SQL query
formatted_query = sqlparse.format(
query,
reindent_aligned=True,
indent_width=2,
# keyword_case='upper' # works poorly in a few cases
)
formatter = TerminalFormatter()
if self.cli_settings.get(
'highlight') and self.cli_settings.get(
'highlight_truecolor'):
formatter = TerminalTrueColorFormatter(
style=CHPygmentsStyle)
print(
'\n' +
pygments.highlight(formatted_query, CHLexer(), formatter))
# TODO: use sqlparse's parser instead
query_split = query.split()
if not query_split:
return Response(query, fmt)
# Since sessions aren't supported over HTTP, we have to make some quirks:
# USE database;
if query_split[0].upper() == 'USE' and len(query_split) == 2:
old_database = self.database
self.database = query_split[1]
try:
self.test_query()
except DBException as e:
self.database = old_database
raise e
return Response(
query,
fmt,
message='Changed the current database to {0}.'.format(
self.database))
# Set response format
if query_split[0].upper(
) in FORMATTABLE_QUERIES and len(query_split) >= 2:
if query_split[-2].upper() == 'FORMAT':
fmt = query_split[-1]
elif query_split[-2].upper() != 'FORMAT':
if query_split[0].upper() != 'INSERT' or data is not None:
if query[-2:] in (r'\g', r'\G'):
query = query[:-2] + ' FORMAT Vertical'
else:
query = query + ' FORMAT {fmt}'.format(fmt=fmt)
params = {
'database': self.database,
'stacktrace': int(self.stacktrace)
}
if query_id:
params['query_id'] = query_id
has_outfile = False
if query_split[0].upper() == 'SELECT':
# Detect INTO OUTFILE at the end of the query
t_query = [
t.value.upper() if t.ttype == Keyword else t.value
for t in sqlparse.parse(query)[0]
if t.ttype not in (Whitespace, Newline)
]
try:
last_tokens = t_query[-5:]
into_pos = last_tokens.index('INTO')
has_outfile = into_pos >= 0 and last_tokens.index(
'OUTFILE') == into_pos + 1
if has_outfile:
path = last_tokens[into_pos + 2].strip("'")
# Remove `INTO OUTFILE '/path/to/file.out'`
last_tokens.pop(into_pos)
last_tokens.pop(into_pos)
last_tokens.pop(into_pos)
query = ' '.join(t_query[:-5] + last_tokens)
except ValueError:
has_outfile = False
method = 'POST'
response = self._query(method,
query,
params,
fmt=fmt,
stream=stream,
data=data,
compress=compress,
**kwargs)
if has_outfile:
try:
with open(path, 'wb') as f:
if not f:
return response
if stream:
for line in response.iter_lines():
f.write(line)
else:
f.write(response.data.encode())
except Exception as e:
echo.warning(
"Caught an exception when writing to file: {0}".format(e))
return response
def query(self, query, data=None, fmt='PrettyCompactMonoBlock', stream=False, verbose=False, query_id=None, **kwargs):
query = sqlparse.format(query, strip_comments=True).rstrip(';')
if kwargs.pop('show_formatted', False) and verbose:
# Highlight & reformat the SQL query
formatted_query = sqlparse.format(
query,
reindent=True,
indent_width=4,
# keyword_case='upper' # works poorly in a few cases
)
print('\n' + pygments.highlight(
formatted_query,
CHLexer(),
TerminalTrueColorFormatter(style=CHPygmentsStyle)
))
# TODO: use sqlparse's parser instead
query_split = query.split()
if len(query_split) == 0:
return Response(query, fmt)
# Since sessions aren't supported over HTTP, we have to make some quirks:
# USE database;
if query_split[0].upper() == 'USE' and len(query_split) == 2:
old_database = self.database
self.database = query_split[1]
try:
self.test_query()
except DBException as e:
self.database = old_database
raise e
return Response(query, fmt, message='Changed the current database to {0}.'.format(self.database))
# SET foo = bar;
if query_split[0].upper() == 'SET' and len(query_split) == 4:
key, value = query_split[1], query_split[3]
old_value = self.settings.get(key)
self.settings[key] = value
try:
self.test_query()
except DBException as e:
if old_value is not None:
self.settings[key] = old_value
else:
del self.settings[key]
raise e
return Response(query, fmt, response='', message='Set {0} to {1}.'.format(key, value))
# Set response format
if query_split[0].upper() in FORMATTABLE_QUERIES and len(query_split) >= 2:
if query_split[-2].upper() == 'FORMAT':
fmt = query_split[-1]
elif query_split[-2].upper() != 'FORMAT':
if query_split[0].upper() != 'INSERT' or data is not None:
if query.endswith('\G') or query.endswith('\g'):
query = query[:-2] + ' FORMAT Vertical'
else:
query = query + ' FORMAT {fmt}'.format(fmt=fmt)
params = {'database': self.database, 'stacktrace': int(self.stacktrace)}
if query_id:
params['query_id'] = query_id
params.update(self.settings)
# Detect INTO OUTFILE at the end of the query
t_query = [
t.value.upper() if t.ttype == Keyword else t.value
for t in sqlparse.parse(query)[0]
if t.ttype not in (Whitespace, Newline)
]
try:
last_tokens = t_query[-5:]
into_pos = last_tokens.index('INTO')
has_outfile = into_pos >= 0 and last_tokens.index('OUTFILE') == into_pos + 1
if has_outfile:
path = last_tokens[into_pos + 2].strip("'")
# Remove `INTO OUTFILE '/path/to/file.out'`
last_tokens.pop(into_pos)
last_tokens.pop(into_pos)
last_tokens.pop(into_pos)
query = ' '.join(t_query[:-5] + last_tokens)
except ValueError:
has_outfile = False
if data is not None:
data = ''.join(data)
response = self._query(query, params, fmt=fmt, stream=stream, data=data, **kwargs)
if has_outfile:
try:
with open(path, 'wb') as f:
if not f:
return response
if stream:
for line in response.data:
f.write(line)
f.write(b'\n')
else:
f.write(response.data.encode())
except Exception as e:
echo.warning("Caught an exception when writing to file: {0}".format(e))
return response