Skip to content

Commit

Permalink
Fixed proper encoding when opening file, limiting zipping from stdin
Browse files Browse the repository at this point in the history
  • Loading branch information
harelba committed Oct 30, 2014
1 parent 8d45a4a commit 5043992
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 6 deletions.
22 changes: 16 additions & 6 deletions bin/q
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,11 @@ class BadHeaderException(Exception):
return repr(self.msg)


class CannotUnzipStdInException(Exception):

def __init__(self):
pass

class EmptyDataException(Exception):

def __init__(self):
Expand Down Expand Up @@ -714,13 +719,15 @@ class TableCreator(object):

# Check if it's standard input or a file
if filename == '-':
f = sys.stdin
codec_info = codecs.lookup(self.encoding)
f = codecs.StreamReaderWriter(sys.stdin,codec_info.streamreader,codec_info.streamwriter,None)
if self.gzipped:
raise CannotUnzipStdInException()
else:
f = file(filename, 'rb')

# Wrap it with gzip decompression if needed
if self.gzipped or filename.endswith('.gz'):
f = gzip.GzipFile(fileobj=f)
if self.gzipped or filename.endswith('.gz'):
f = gzip.GzipFile(fileobj=file(filename,'rb'))
else:
f = codecs.open(filename, 'rb',encoding=self.encoding)

self.read_file_using_csv(f, analyze_only)
if not self.table_created:
Expand Down Expand Up @@ -1054,6 +1061,9 @@ except (UnicodeDecodeError, UnicodeError), e:
except BadHeaderException, e:
print >>sys.stderr, "Bad header row: %s" % e.msg
sys.exit(35)
except CannotUnzipStdInException,e:
print >>sys.stderr,"Cannot decompress standard input. Pipe the input through zcat in order to decompress."
sys.exit(36)
except KeyboardInterrupt:
print >>sys.stderr, "Interrupted"
sys.exit(0)
Expand Down
15 changes: 15 additions & 0 deletions test/test-suite
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,21 @@ class BasicTests(AbstractQTestCase):

self.cleanup(tmpfile)

def test_attempt_to_unzip_stdin(self):
tmpfile = self.create_file_with_data(
'\x1f\x8b\x08\x08\xf2\x18\x12S\x00\x03xxxxxx\x003\xe42\xe22\xe62\xe12\xe52\xe32\xe7\xb2\xe0\xb2\xe424\xe0\x02\x00\xeb\xbf\x8a\x13\x15\x00\x00\x00', encoding='none')

cmd = 'cat %s | ../bin/q -z "select sum(c1),avg(c1) from -"' % tmpfile.name

retcode, o, e = run_command(cmd)
self.assertTrue(retcode != 0)
self.assertTrue(len(o) == 0)
self.assertTrue(len(e) == 1)

self.assertEquals(e[0],'Cannot decompress standard input. Pipe the input through zcat in order to decompress.')

self.cleanup(tmpfile)

def test_delimition_mistake_with_header(self):
tmpfile = self.create_file_with_data(sample_data_no_header)

Expand Down

0 comments on commit 5043992

Please sign in to comment.