Support multiple result sets (#130)

* Add multiple result set support to mssqlcliclient

* Add multi result set support

* Add tests for multi statement and multi result sets

Author: abhisheksinha89

mssqlcli/mssqlcliclient.py

@@ -121,7 +121,8 @@ def execute_multi_statement_single_batch(self, query):
                 if not sql:
                     yield None, None, None, sql, False
                     continue
-                yield self.execute_single_batch_query(sql)
+                for rows, columns, status, sql, is_error in self.execute_single_batch_query(sql):
+                    yield rows, columns, status, sql, is_error
 
     def execute_single_batch_query(self, query):
         if not self.is_connected:
@@ -134,22 +135,22 @@ def execute_single_batch_query(self, query):
         query_request = self.sql_tools_client.create_request(u'query_execute_string_request',
                                                              {u'OwnerUri': self.owner_uri, u'Query': query})
         query_request.execute()
-        query_message = None
+        query_messages = []
         while not query_request.completed():
             query_response = query_request.get_response()
             if query_response:
                 if isinstance(query_response, queryservice.QueryExecuteErrorResponseEvent):
-                    return self.tabular_results_generator(column_info=None, result_rows=None,
+                    yield self.tabular_results_generator(column_info=None, result_rows=None,
                                                           query=query, message=query_response.error_message,
                                                           is_error=True)
                 elif isinstance(query_response, queryservice.QueryMessageEvent):
-                    query_message = query_response
+                    query_messages.append(query_response)
             else:
                 sleep(time_wait_if_no_response)
 
         if query_response.exception_message:
             logger.error(u'Query response had an exception')
-            return self.tabular_results_generator(
+            yield self.tabular_results_generator(
                 column_info=None,
                 result_rows=None,
                 query=query,
@@ -158,42 +159,43 @@ def execute_single_batch_query(self, query):
 
         if (not query_response.batch_summaries[0].result_set_summaries) or \
            (query_response.batch_summaries[0].result_set_summaries[0].row_count == 0):
-            return self.tabular_results_generator(
+            yield self.tabular_results_generator(
                 column_info=None,
                 result_rows=None,
                 query=query,
-                message=query_message.message if query_message else u'',
-                is_error=query_message.is_error if query_message else query_response.batch_summaries[0].has_error)
+                message=query_messages[0].message if query_messages else u'',
+                is_error=query_messages[0].is_error if query_messages else query_response.batch_summaries[0].has_error)
 
         else:
-            query_subset_request = self.sql_tools_client.create_request(u'query_subset_request',
-                                                                        {u'OwnerUri': query_response.owner_uri,
-                                                                         u'BatchIndex': query_response.batch_summaries[0].result_set_summaries[0].batch_id,
-                                                                            u'ResultSetIndex': query_response.batch_summaries[0].result_set_summaries[0].id,
-                                                                            u'RowsStartIndex': 0,
-                                                                            u'RowCount': query_response.batch_summaries[0].result_set_summaries[0].row_count})
-
-            query_subset_request.execute()
-            while not query_subset_request.completed():
-                subset_response = query_subset_request.get_response()
-                if not subset_response:
-                    sleep(time_wait_if_no_response)
-
-            if subset_response.error_message:
-                logger.error(u'Error obtaining result sets for query response')
-                return self.tabular_results_generator(
-                    column_info=None,
-                    result_rows=None,
+            for result_set_summary in query_response.batch_summaries[0].result_set_summaries:
+                query_subset_request = self.sql_tools_client.create_request(u'query_subset_request',
+                                                                            {u'OwnerUri': query_response.owner_uri,
+                                                                             u'BatchIndex': result_set_summary.batch_id,
+                                                                                u'ResultSetIndex': result_set_summary.id,
+                                                                                u'RowsStartIndex': 0,
+                                                                                u'RowCount': result_set_summary.row_count})
+
+                query_subset_request.execute()
+                while not query_subset_request.completed():
+                    subset_response = query_subset_request.get_response()
+                    if not subset_response:
+                        sleep(time_wait_if_no_response)
+
+                if subset_response.error_message:
+                    logger.error(u'Error obtaining result sets for query response')
+                    yield self.tabular_results_generator(
+                        column_info=None,
+                        result_rows=None,
+                        query=query,
+                        message=subset_response.error_message,
+                        is_error=True)
+
+                logger.info(u'Obtained result set for query')
+                yield self.tabular_results_generator(
+                    column_info=result_set_summary.column_info,
+                    result_rows=subset_response.rows,
                     query=query,
-                    message=subset_response.error_message,
-                    is_error=True)
-
-            logger.info(u'Obtained result set for query')
-            return self.tabular_results_generator(
-                column_info=query_response.batch_summaries[0].result_set_summaries[0].column_info,
-                result_rows=subset_response.rows,
-                query=query,
-                message=query_message.message if query_message else u'')
+                    message=query_messages[result_set_summary.id].message if query_messages else u'')
 
     def tabular_results_generator(
             self, column_info, result_rows, query, message, is_error=False):
@@ -213,55 +215,63 @@ def schemas(self):
         """ Returns a list of schema names"""
         query = mssqlqueries.get_schemas()
         logger.info(u'Schemas query: {0}'.format(query))
-        return [x[0] for x in self.execute_single_batch_query(query)[0]]
+        for tabular_result in self.execute_single_batch_query(query):
+            return [x[0] for x in tabular_result[0]]
 
     def databases(self):
         """ Returns a list of database names"""
         query = mssqlqueries.get_databases()
         logger.info(u'Databases query: {0}'.format(query))
-        return [x[0] for x in self.execute_single_batch_query(query)[0]]
+        for tabular_result in self.execute_single_batch_query(query):
+            return [x[0] for x in tabular_result[0]]
 
     def tables(self):
         """ Yields (schema_name, table_name) tuples"""
         query = mssqlqueries.get_tables()
         logger.info(u'Tables query: {0}'.format(query))
-        for row in self.execute_single_batch_query(query)[0]:
-            yield (row[0], row[1])
+        for tabular_result in self.execute_single_batch_query(query):
+            for row in tabular_result[0]:
+                yield (row[0], row[1])
 
     def table_columns(self):
         """ Yields (schema_name, table_name, column_name, data_type, column_default) tuples"""
         query = mssqlqueries.get_table_columns()
         logger.info(u'Table columns query: {0}'.format(query))
-        for row in self.execute_single_batch_query(query)[0]:
-            yield (row[0], row[1], row[2], row[3], row[4])
+        for tabular_result in self.execute_single_batch_query(query):
+            for row in tabular_result[0]:
+                yield (row[0], row[1], row[2], row[3], row[4])
 
     def views(self):
         """ Yields (schema_name, table_name) tuples"""
         query = mssqlqueries.get_views()
         logger.info(u'Views query: {0}'.format(query))
-        for row in self.execute_single_batch_query(query)[0]:
-            yield (row[0], row[1])
+        for tabular_result in self.execute_single_batch_query(query):
+            for row in tabular_result[0]:
+                yield (row[0], row[1])
 
     def view_columns(self):
         """ Yields (schema_name, table_name, column_name, data_type, column_default) tuples"""
         query = mssqlqueries.get_view_columns()
         logger.info(u'View columns query: {0}'.format(query))
-        for row in self.execute_single_batch_query(query)[0]:
-            yield (row[0], row[1], row[2], row[3], row[4])
+        for tabular_result in self.execute_single_batch_query(query):
+            for row in tabular_result[0]:
+                yield (row[0], row[1], row[2], row[3], row[4])
 
     def user_defined_types(self):
         """ Yields (schema_name, type_name) tuples"""
         query = mssqlqueries.get_user_defined_types()
         logger.info(u'UDTs query: {0}'.format(query))
-        for row in self.execute_single_batch_query(query)[0]:
-            yield (row[0], row[1])
+        for tabular_result in self.execute_single_batch_query(query):
+            for row in tabular_result[0]:
+                yield (row[0], row[1])
 
     def foreignkeys(self):
         """ Yields (parent_schema, parent_table, parent_column, child_schema, child_table, child_column) typles"""
         query = mssqlqueries.get_foreignkeys()
         logger.info(u'Foreign keys query: {0}'.format(query))
-        for row in self.execute_single_batch_query(query)[0]:
-            yield ForeignKey(*row)
+        for tabular_result in self.execute_single_batch_query(query):
+            for row in tabular_result[0]:
+                yield ForeignKey(*row)
 
     def shutdown(self):
         self.sql_tools_client.shutdown()

tests/mssqlutils.py

@@ -51,14 +51,14 @@ def run_and_return_string_from_formatter(client, sql, join=False, expanded=False
     :return:
     """
 
-    rows, col, message, query, is_error = client.execute_single_batch_query(sql)
-    settings = OutputSettings(table_format='psql', dcmlfmt='d', floatfmt='g',
-                              expanded=expanded)
-    formatted = format_output(None, rows, col, message, settings)
-    if join:
-        formatted = '\n'.join(formatted)
+    for rows, col, message, query, is_error in client.execute_single_batch_query(sql):
+        settings = OutputSettings(table_format='psql', dcmlfmt='d', floatfmt='g',
+                                  expanded=expanded)
+        formatted = format_output(None, rows, col, message, settings)
+        if join:
+            formatted = '\n'.join(formatted)
 
-    return formatted
+        return formatted
 
 
 def shutdown(connection):

tests/test_mssqlcliclient.py

@@ -88,10 +88,9 @@ def test_get_query_results(self):
                 select 3, 'Night'
             """
 
-            rows, col, message, query, is_error = client.execute_single_batch_query(test_query)
-
-            self.assertTrue(len(rows), 3)
-            self.assertTrue(query, test_query)
+            for rows, col, message, query, is_error in client.execute_single_batch_query(test_query):
+                self.assertTrue(len(rows), 3)
+                self.assertTrue(query, test_query)
         finally:
             shutdown(client)
 
@@ -118,11 +117,11 @@ def test_schema_table_views_and_columns_query(self):
         """
         try:
             client = create_mssql_cli_client()
-            client.execute_single_batch_query('CREATE TABLE tabletest1 (a int, b varchar(25));')
-            client.execute_single_batch_query('CREATE TABLE tabletest2 (x int, y varchar(25), z bit);')
-            client.execute_single_batch_query('CREATE VIEW viewtest as SELECT a from tabletest1;')
-            client.execute_single_batch_query('CREATE SCHEMA schematest;')
-            client.execute_single_batch_query('CREATE TABLE schematest.tabletest1 (a int);')
+            list(client.execute_single_batch_query('CREATE TABLE tabletest1 (a int, b varchar(25));'))
+            list(client.execute_single_batch_query('CREATE TABLE tabletest2 (x int, y varchar(25), z bit);'))
+            list(client.execute_single_batch_query('CREATE VIEW viewtest as SELECT a from tabletest1;'))
+            list(client.execute_single_batch_query('CREATE SCHEMA schematest;'))
+            list(client.execute_single_batch_query('CREATE TABLE schematest.tabletest1 (a int);'))
 
             assert ('schematest', 'tabletest1') in set(client.tables())
             assert ('dbo', 'viewtest') in set(client.views())
@@ -131,11 +130,11 @@ def test_schema_table_views_and_columns_query(self):
             assert 'schematest' in client.schemas()
 
         finally:
-            client.execute_single_batch_query('DROP TABLE tabletest1;')
-            client.execute_single_batch_query('DROP TABLE tabletest2;')
-            client.execute_single_batch_query('DROP VIEW viewtest IF EXISTS;')
-            client.execute_single_batch_query('DROP TABLE schematest.tabletest1;')
-            client.execute_single_batch_query('DROP SCHEMA schematest;')
+            list(client.execute_single_batch_query('DROP TABLE tabletest1;'))
+            list(client.execute_single_batch_query('DROP TABLE tabletest2;'))
+            list(client.execute_single_batch_query('DROP VIEW viewtest IF EXISTS;'))
+            list(client.execute_single_batch_query('DROP TABLE schematest.tabletest1;'))
+            list(client.execute_single_batch_query('DROP SCHEMA schematest;'))
             shutdown(client)
 
     def test_mssqlcliclient_reset_connection(self):
@@ -155,12 +154,54 @@ def test_mssqlcliclient_multiple_statement(self):
         try:
             client = create_mssql_cli_client()
             multi_statement_query = u"select 'Morning' as [Name] UNION ALL select 'Evening'; select 1;"
+            multi_statement_query2 = u"select 1; select 'foo' from teapot;"
+            multi_statement_query3 = u"select 'foo' from teapot; select 2;"
             for rows, col, message, query, is_error in \
                 client.execute_multi_statement_single_batch(multi_statement_query):
                 if query == u"select 'Morning' as [Name] UNION ALL select 'Evening'":
                     self.assertTrue(len(rows), 2)
                 else:
                     self.assertTrue(len(rows), 1)
+
+            for rows, col, message, query, is_error in \
+                    client.execute_multi_statement_single_batch(multi_statement_query2):
+                if query == u"select 1":
+                    self.assertTrue(len(rows) == 1)
+                else:
+                    self.assertTrue(is_error)
+
+            for rows, col, message, query, is_error in \
+                    client.execute_multi_statement_single_batch(multi_statement_query3):
+                if query == u"select 2":
+                    self.assertTrue(len(rows) == 1)
+                else:
+                    self.assertTrue(is_error)
+
+        finally:
+            shutdown(client)
+
+    def test_stored_proc_multiple_result_sets(self):
+        """
+            Verify the results of running a stored proc with multiple result sets
+        """
+        try:
+            client = create_mssql_cli_client()
+            create_stored_proc = u"CREATE PROC sp_mssqlcli_multiple_results " \
+                          u"AS " \
+                          u"BEGIN " \
+                          u"SELECT 'Morning' as [Name] UNION ALL select 'Evening' " \
+                          u"SELECT 'Dawn' as [Name] UNION ALL select 'Dusk' UNION ALL select 'Midnight' " \
+                          u"END"
+            exec_stored_proc = u"EXEC sp_mssqlcli_multiple_results"
+            del_stored_proc = u"DROP PROCEDURE sp_mssqlcli_multiple_results"
+
+            list(client.execute_single_batch_query(create_stored_proc))
+            row_counts = []
+            for rows, columns, message, query, is_error in client.execute_single_batch_query(exec_stored_proc):
+                row_counts.append(len(rows))
+            self.assertTrue(row_counts[0] == 2)
+            self.assertTrue(row_counts[1] == 3)
+            list(client.execute_single_batch_query(del_stored_proc))
         finally:
             shutdown(client)