Trying to fix incorrect set_output_fields on Select.

python-bonobo · Apr 1, 2018 · 521865f · 521865f
1 parent dc2be38
commit 521865f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/bonobo_sqlalchemy/readers.py b/bonobo_sqlalchemy/readers.py
@@ -52,11 +52,15 @@ def __call__(self, context, *, engine):
         query = self.query.strip(' \n;')
 
         offset = 0
+        max_limit = self.limit or self.pack_size
+
+        assert self.pack_size > 0, 'Pack size must be > 0 for now.'
+
         while not self.limit or offset * self.pack_size < self.limit:
             results = engine.execute(
                 '{query} LIMIT {limit}{offset}'.format(
                     query=query,
-                    limit=self.pack_size,
+                    limit=max(min(self.pack_size, max_limit - offset * self.pack_size), 0),
                     offset=' OFFSET {}'.format(offset * self.pack_size) if offset else ''
                 ),
                 use_labels=True
@@ -65,8 +69,8 @@ def __call__(self, context, *, engine):
             if not len(results):
                 break
 
-            for i, row in enumerate(results):
-                if not i:
+            for row in results:
+                if not context.output_type:
                     context.set_output_fields(row.keys())
                 yield tuple(row)
 

diff --git a/examples/sequel_extract.py b/examples/sequel_extract.py
@@ -16,7 +16,7 @@ def extract(context):
 def get_graph(**options):
     graph = bonobo.Graph()
     graph.add_chain(
-        bonobo_sqlalchemy.Select('SELECT * FROM example', limit=100),
+        bonobo_sqlalchemy.Select('SELECT * FROM example', limit=100, pack_size=9),
         bonobo.PrettyPrinter(),
     )