#3_#85: adding to query table_name property and pretty_render function (#86)

Co-authored-by: Netanel Zakay <nezakay@microsoft.com>
Co-authored-by: Yonatan Most <>

pykusto/query.py

@@ -3,6 +3,7 @@ from copy import copy, deepcopy
 from itertools import chain
 from types import FunctionType
 from typing import Tuple, List, Union, Optional
+from os import linesep
 
 from pykusto.client import Table, KustoResponse
 from pykusto.enums import Order, Nulls, JoinKind, Distribution, BagExpansion
@@ -19,10 +20,12 @@ from pykusto.udf import stringify_python_func
 class Query:
     _head: Optional['Query']
     _table: Optional[Table]
+    _table_name: Optional[str]
 
     def __init__(self, head=None) -> None:
         self._head = head if isinstance(head, Query) else None
         self._table = head if isinstance(head, Table) else None
+        self._table_name = head if isinstance(head, str) else None
 
     def __add__(self, other: 'Query') -> 'Query':
         self_copy = deepcopy(self)
@@ -140,7 +143,10 @@ class Query:
     def _compile_all(self, use_full_table_name) -> KQL:
         if self._head is None:
             if self._table is None:
-                return KQL("")
+                if self._table_name is None:
+                    return KQL("")
+                else:
+                    return KQL(self.get_table_name())
             else:
                 table = self._table
                 if use_full_table_name:
@@ -156,11 +162,23 @@ class Query:
         else:
             return self._head.get_table()
 
+    def get_table_name(self) -> str:
+        if self._head is None:
+            return self._table_name
+        else:
+            return self._head.get_table_name()
+
     def render(self, use_full_table_name: bool = False) -> KQL:
         result = self._compile_all(use_full_table_name)
         logger.debug("Complied query: " + result)
         return result
 
+    def pretty_render(self, use_full_table_name: bool = False) -> KQL:
+        kql = self.render(use_full_table_name)
+        if kql is not None:
+            kql = KQL(kql.replace(" |", linesep + "|"))
+        return kql
+
     def execute(self, table: Table = None) -> KustoResponse:
         if self.get_table() is None:
             if table is None:

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='pykusto',
-    version='0.0.20',
+    version='0.0.21',
     packages=find_packages(exclude=['test']),
     url='https://github.com/Azure/pykusto',
     license='MIT License',

test/test_query.py

@@ -1,3 +1,5 @@
+from os import linesep
+
 import pandas as pd
 
 from pykusto.client import PyKustoClient
@@ -64,6 +66,47 @@ class TestQuery(TestBase):
             query_b.render(),
         )
 
+    def test_add_queries_with_table_name(self):
+        query_a = Query('test_table').where(col.numField > 4)
+        query_b = Query().take(5)
+        query = query_a + query_b
+        self.assertEqual(
+            "test_table | where numField > 4 | take 5",
+            query.render(),
+        )
+        self.assertEqual(
+            "test_table",
+            query.get_table_name(),
+        )
+
+        # make sure the originals didn't change
+        self.assertEqual(
+            "test_table | where numField > 4",
+            query_a.render(),
+        )
+        self.assertEqual(
+            "test_table",
+            query_a.get_table_name(),
+        )
+
+        self.assertEqual(
+            " | take 5",
+            query_b.render(),
+        )
+        self.assertEqual(
+            None,
+            query_b.get_table_name(),
+        )
+
+    def test_pretty_render(self):
+        query = Query('test_table').where(col.numField > 4).take(5)
+        self.assertEqual(
+            "test_table" + linesep +
+            "| where numField > 4" + linesep +
+            "| take 5",
+            query.pretty_render(),
+        )
+
     def test_where(self):
         self.assertEqual(
             "test_table | where numField > 4",
@@ -223,7 +266,9 @@ class TestQuery(TestBase):
     def test_mv_expand_assign_to_with_assign_other_params(self):
         self.assertEqual(
             "test_table | mv-expand bagexpansion=bag with_itemindex=foo expanded_field = arrayField, expanded_field2 = arrayField2 limit 4",
-            Query(t).mv_expand(t.arrayField.assign_to(col.expanded_field), expanded_field2=t.arrayField2, bag_expansion=BagExpansion.BAG, with_item_index=col.foo, limit=4).render(),
+            Query(t).mv_expand(
+                t.arrayField.assign_to(col.expanded_field), expanded_field2=t.arrayField2, bag_expansion=BagExpansion.BAG, with_item_index=col.foo, limit=4
+            ).render(),
         )
 
     def test_mv_expand_assign_multiple(self):