Skip to content

Commit 9bfca3f

Browse files
harelbaharelba
committed
Added option to add output header to the output + Tests
1 parent 471371d commit 9bfca3f

2 files changed

Lines changed: 65 additions & 4 deletions

File tree

‎bin/q‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ default_formatting = get_option_with_default(p, 'string', 'formatting', None)
8787
default_encoding = get_option_with_default(p, 'string', 'encoding', 'UTF-8')
8888
default_output_encoding = get_option_with_default(p, 'string', 'encoding', None)
8989
default_query_encoding = get_option_with_default(p, 'string', 'query_encoding', locale.getpreferredencoding())
90+
default_output_header = get_option_with_default(p, 'string', 'output_header', False)
9091

9192
parser = OptionParser(usage="""
9293
q allows performing SQL-like statements on tabular text data.
@@ -130,6 +131,7 @@ parser.add_option("-T", "--tab-delimited-output", dest="tab_delimited_output", d
130131
help="Same as -D <tab>. Just a shorthand for outputing tab delimited output. You can use -D $'\t' if you want.")
131132
parser.add_option("-H", "--skip-header", dest="skip_header", default=default_skip_header, action="store_true",
132133
help="Skip header row. This has been changed from earlier version - Only one header row is supported, and the header row is used for column naming")
134+
parser.add_option("-O", "--output-header", dest="output_header", default=default_output_header, action="store_true",help="Output header line. Output column-names are determined from the query itself. Use column aliases in order to set your column names in the query. For example, 'select name FirstName,value1/value2 MyCalculation from ...'. This can be used even if there was no header in the input.")
133135
parser.add_option("-f", "--formatting", dest="formatting", default=default_formatting,
134136
help="Output-level formatting, in the format X=fmt,Y=fmt etc, where X,Y are output column numbers (e.g. 1 for first SELECT column etc.")
135137
parser.add_option("-e", "--encoding", dest="encoding", default=default_encoding,
@@ -155,6 +157,10 @@ parser.add_option("-k", "--keep-leading-whitespace", dest="keep_leading_whitespa
155157
def regexp(regular_expression, data):
156158
return re.search(regular_expression, data) is not None
157159

160+
class Sqlite3DBResults(object):
161+
def __init__(self,query_column_names,results):
162+
self.query_column_names = query_column_names
163+
self.results = results
158164

159165
class Sqlite3DB(object):
160166

@@ -184,12 +190,17 @@ class Sqlite3DB(object):
184190
def execute_and_fetch(self, q):
185191
try:
186192
if self.show_sql:
187-
print q
193+
print repr(q)
188194
self.cursor.execute(q)
195+
if self.cursor.description is not None:
196+
# we decode the column names, so they can be encoded to any output format later on
197+
query_column_names = [c[0].decode('utf-8') for c in self.cursor.description]
198+
else:
199+
query_column_names = None
189200
result = self.cursor.fetchall()
190201
finally:
191202
pass # cursor.close()
192-
return result
203+
return Sqlite3DBResults(query_column_names,result)
193204

194205
def _get_as_list_str(self, l):
195206
return ",".join(['"%s"' % x.replace('"', '""') for x in l])
@@ -299,6 +310,8 @@ class Sql(object):
299310
# names
300311
self.qtable_name_effective_table_names = {}
301312

313+
self.query_column_names = None
314+
302315
# Go over all sql parts
303316
idx = 0
304317
while idx < len(self.sql_parts):
@@ -358,7 +371,8 @@ class Sql(object):
358371
return " ".join(effective_sql)
359372

360373
def execute_and_fetch(self, db):
361-
return db.execute_and_fetch(self.get_effective_sql())
374+
db_results_obj = db.execute_and_fetch(self.get_effective_sql())
375+
return db_results_obj
362376

363377

364378
class LineSplitter(object):
@@ -971,7 +985,9 @@ try:
971985
sys.exit(0)
972986

973987
# Execute the query and fetch the data
974-
m = sql_object.execute_and_fetch(db)
988+
db_results_obj = sql_object.execute_and_fetch(db)
989+
m = db_results_obj.results
990+
output_column_name_list = db_results_obj.query_column_names
975991
except EmptyDataException:
976992
print >>sys.stderr, "Warning - data is empty"
977993
sys.exit(0)
@@ -1023,6 +1039,8 @@ else:
10231039
formatting_dict = None
10241040

10251041
try:
1042+
if options.output_header and output_column_name_list is not None:
1043+
m.insert(0,output_column_name_list)
10261044
for rownum, row in enumerate(m):
10271045
row_str = []
10281046
for i, col in enumerate(row):

‎test/test-suite‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,20 @@ class BasicTests(AbstractQTestCase):
330330

331331
self.cleanup(tmpfile)
332332

333+
def test_output_header_when_input_header_exists(self):
334+
tmpfile = self.create_file_with_data(sample_data_with_header)
335+
cmd = '../bin/q -d , "select name from %s" -H -O' % tmpfile.name
336+
retcode, o, e = run_command(cmd)
337+
338+
self.assertEquals(retcode, 0)
339+
self.assertEquals(len(o), 4)
340+
self.assertEquals(o[0],'name')
341+
self.assertEquals(o[1],'a')
342+
self.assertEquals(o[2],'b')
343+
self.assertEquals(o[3],'c')
344+
345+
self.cleanup(tmpfile)
346+
333347
def test_generated_column_name_warning_when_header_line_exists(self):
334348
tmpfile = self.create_file_with_data(sample_data_with_header)
335349
cmd = '../bin/q -d , "select c3 from %s" -H' % tmpfile.name
@@ -558,6 +572,25 @@ class BasicTests(AbstractQTestCase):
558572
self.cleanup(tmp_data_file)
559573
self.cleanup(tmp_query_file)
560574

575+
def test_output_header_with_non_ascii_names(self):
576+
tmp_data_file = self.create_file_with_data(sample_data_with_header)
577+
tmp_query_file = self.create_file_with_data("select name,'Hr\xc3\xa1\xc4\x8d' Hr\xc3\xa1\xc4\x8d from %s" % tmp_data_file.name,encoding=None)
578+
579+
cmd = '../bin/q -d , -q %s -H -Q utf-8 -O' % tmp_query_file.name
580+
retcode, o, e = run_command(cmd)
581+
582+
self.assertEquals(retcode,0)
583+
self.assertEquals(len(o),4)
584+
self.assertEquals(len(e),0)
585+
586+
self.assertEquals(o[0].decode(SYSTEM_ENCODING), u'name,Hr\xe1\u010d')
587+
self.assertEquals(o[1].decode(SYSTEM_ENCODING), u'a,Hr\xe1\u010d')
588+
self.assertEquals(o[2].decode(SYSTEM_ENCODING), u'b,Hr\xe1\u010d')
589+
self.assertEquals(o[3].decode(SYSTEM_ENCODING), u'c,Hr\xe1\u010d')
590+
591+
self.cleanup(tmp_data_file)
592+
self.cleanup(tmp_query_file)
593+
561594
def test_use_query_file_with_query_encoding(self):
562595
tmp_data_file = self.create_file_with_data(sample_data_with_header)
563596
tmp_query_file = self.create_file_with_data("select name,'Hr\xc3\xa1\xc4\x8d' from %s" % tmp_data_file.name,encoding=None)
@@ -907,6 +940,16 @@ class SqlTests(AbstractQTestCase):
907940
self.assertEquals(o[0], 'ppp dip.1@otherdomain.com')
908941
self.assertEquals(o[1], 'ppp dip.2@otherdomain.com')
909942

943+
def test_join_example_with_output_header(self):
944+
cmd = '../bin/q -O "select myfiles.c8 aaa,emails.c2 bbb from ../examples/exampledatafile myfiles join ../examples/group-emails-example emails on (myfiles.c4 = emails.c1) where myfiles.c8 = \'ppp\'"'
945+
retcode, o, e = run_command(cmd)
946+
947+
self.assertEquals(retcode, 0)
948+
self.assertEquals(len(o), 3)
949+
950+
self.assertEquals(o[0], 'aaa bbb')
951+
self.assertEquals(o[1], 'ppp dip.1@otherdomain.com')
952+
self.assertEquals(o[2], 'ppp dip.2@otherdomain.com')
910953

911954
def suite():
912955
tl = unittest.TestLoader()

0 commit comments

Comments
 (0)