diff --git a/pykusto/query.py b/pykusto/query.py index c31ebaa..6427da2 100644 --- a/pykusto/query.py +++ b/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: diff --git a/setup.py b/setup.py index 212c442..5f452e5 100644 --- a/setup.py +++ b/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', diff --git a/test/test_query.py b/test/test_query.py index 4473972..a519e5a 100644 --- a/test/test_query.py +++ b/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):