Repository: Azure/pykusto

Issue:
Add support for cross-cluster join
If the joined table is from a different cluster, render it using the following syntax: cluster("a").database("b").table("c")

Comments:

Patch:
diff --git a/pykusto/client.py b/pykusto/client.py
index 1d708c7..2bc47e4 100644
--- a/pykusto/client.py
+++ b/pykusto/client.py
@@ -1,5 +1,8 @@
 from typing import Union, List, Tuple
 
+# noinspection PyProtectedMember
+from urllib.parse import urlparse
+
 # noinspection PyProtectedMember
 from azure.kusto.data._response import KustoResponseDataSet
 from azure.kusto.data.request import KustoClient, KustoConnectionStringBuilder, ClientRequestProperties
@@ -12,6 +15,7 @@ class PyKustoClient:
     Handle to a Kusto cluster
     """
     _client: KustoClient
+    _cluster_name: str
 
     def __init__(self, client_or_cluster: Union[str, KustoClient]) -> None:
         """
@@ -22,8 +26,11 @@ class PyKustoClient:
         """
         if isinstance(client_or_cluster, KustoClient):
             self._client = client_or_cluster
+            # noinspection PyProtectedMember
+            self._cluster_name = urlparse(client_or_cluster._query_endpoint).netloc  # TODO neater way
         else:
             self._client = self._get_client_for_cluster(client_or_cluster)
+            self._cluster_name = client_or_cluster
 
     def execute(self, database: str, query: KQL, properties: ClientRequestProperties = None) -> KustoResponseDataSet:
         return self._client.execute(database, query, properties)
@@ -35,6 +42,9 @@ class PyKustoClient:
     def __getitem__(self, database_name: str) -> 'Database':
         return Database(self, database_name)
 
+    def get_cluster_name(self) -> str:
+        return self._cluster_name
+
     @staticmethod
     def _get_client_for_cluster(cluster: str) -> KustoClient:
         return KustoClient(KustoConnectionStringBuilder.with_aad_device_authentication(cluster))
@@ -70,7 +80,7 @@ class Table:
     Handle to a Kusto table
     """
     database: Database
-    table: KQL
+    tables: Union[str, List[str], Tuple[str, ...]]
 
     def __init__(self, database: Database, tables: Union[str, List[str], Tuple[str, ...]]) -> None:
         """
@@ -82,19 +92,31 @@ class Table:
         """
 
         self.database = database
-
-        if isinstance(tables, (List, Tuple)):
-            self.table = KQL(', '.join(tables))
+        self.tables = [tables] if isinstance(tables, str) else tables
+
+    def get_table(self) -> KQL:
+        result = KQL(', '.join(self.tables))
+        if '*' in result or ',' in result:
+            result = KQL('union ' + result)
+        return result
+
+    def get_full_table(self) -> KQL:
+        assert len(self.tables) > 0
+        if len(self.tables) == 1 and not any('*' in t for t in self.tables):
+            return self._format_full_table_name(self.tables[0])
         else:
-            self.table = KQL(tables)
-        if '*' in self.table or ',' in self.table:
-            self.table = KQL('union ' + self.table)
+            return KQL("union " + ", ".join(self._format_full_table_name(t) for t in self.tables))
+
+    def _format_full_table_name(self, table):
+        table_format_str = 'cluster("{}").database("{}").table("{}")'
+        return KQL(
+            table_format_str.format(self.database.client.get_cluster_name(), self.database.name, table))
 
     def execute(self, rendered_query: KQL) -> KustoResponseDataSet:
         return self.database.execute(rendered_query)
 
     def show_columns(self) -> Tuple[Tuple[str, str], ...]:
-        res: KustoResponseDataSet = self.execute(KQL('.show table {}'.format(self.table)))
+        res: KustoResponseDataSet = self.execute(KQL('.show table {}'.format(self.get_table())))
         return tuple(
             (
                 r[0],  # Column name
diff --git a/pykusto/query.py b/pykusto/query.py
index 547c1c9..42caba4 100644
--- a/pykusto/query.py
+++ b/pykusto/query.py
@@ -156,14 +156,18 @@ class Query:
     def _compile(self) -> KQL:
         pass
 
-    def _compile_all(self) -> KQL:
+    def _compile_all(self, use_full_table_name) -> KQL:
         if self._head is None:
             if self._table is None:
                 return KQL("")
             else:
-                return self._table.table
+                table = self._table
+                if use_full_table_name:
+                    return table.get_full_table()
+                else:
+                    return table.get_table()
         else:
-            return KQL("{} | {}".format(self._head._compile_all(), self._compile()))
+            return KQL("{} | {}".format(self._head._compile_all(use_full_table_name), self._compile()))
 
     def get_table(self) -> Table:
         if self._head is None:
@@ -171,8 +175,8 @@ class Query:
         else:
             return self._head.get_table()
 
-    def render(self) -> KQL:
-        result = self._compile_all()
+    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
 
@@ -180,7 +184,7 @@ class Query:
         if self.get_table() is None:
             if table is None:
                 raise RuntimeError("No table supplied")
-            rendered_query = table.table + self.render()
+            rendered_query = table.get_table() + self.render()
         else:
             if table is not None:
                 raise RuntimeError("This table is already bound to a query")
@@ -412,7 +416,7 @@ class JoinQuery(Query):
 
         return KQL("join {} ({}) on {}".format(
             "" if self._kind is None else "kind={}".format(self._kind.value),
-            self._joined_query.render(),
+            self._joined_query.render(use_full_table_name=True),
             ", ".join([self._compile_on_attribute(attr) for attr in self._on_attributes])))
 
 


Test patch:
diff --git a/test/test_query.py b/test/test_query.py
index b4ed016..a228e03 100644
--- a/test/test_query.py
+++ b/test/test_query.py
@@ -69,7 +69,7 @@ class TestQuery(TestBase):
         table = PyKustoClient(mock_kusto_client)['test_db']['test_table']
 
         self.assertEqual(
-            " | where foo > 4 | take 5 | join kind=inner (test_table) on col0, $left.col1==$right.col2",
+            ' | where foo > 4 | take 5 | join kind=inner (cluster("test_cluster.kusto.windows.net").database("test_db").table("test_table")) on col0, $left.col1==$right.col2',
             Query().where(col.foo > 4).take(5).join(
                 Query(table), kind=JoinKind.INNER).on(col.col0).on(col.col1, col.col2).render(),
         )
@@ -79,7 +79,7 @@ class TestQuery(TestBase):
         table = PyKustoClient(mock_kusto_client)['test_db']['test_table']
 
         self.assertEqual(
-            " | where foo > 4 | take 5 | join kind=inner (test_table | where bla == 2 | take 6) on col0, "
+            ' | where foo > 4 | take 5 | join kind=inner (cluster("test_cluster.kusto.windows.net").database("test_db").table("test_table") | where bla == 2 | take 6) on col0, '
             "$left.col1==$right.col2",
             Query().where(col.foo > 4).take(5).join(
                 Query(table).where(col.bla == 2).take(6), kind=JoinKind.INNER).on(col.col0).on(col.col1,
diff --git a/test/test_table.py b/test/test_table.py
index d7a81b6..755acb2 100644
--- a/test/test_table.py
+++ b/test/test_table.py
@@ -1,8 +1,10 @@
 from typing import List, Tuple
 from unittest.mock import patch
+from urllib.parse import urljoin
 
 from azure.kusto.data.request import KustoClient, ClientRequestProperties
 
+from pykusto.expressions import column_generator as col
 from pykusto.client import PyKustoClient
 from pykusto.query import Query
 from test.test_base import TestBase
@@ -12,8 +14,9 @@ from test.test_base import TestBase
 class MockKustoClient(KustoClient):
     executions: List[Tuple[str, str, ClientRequestProperties]]
 
-    def __init__(self):
+    def __init__(self, cluster="https://test_cluster.kusto.windows.net"):
         self.executions = []
+        self._query_endpoint = urljoin(cluster, "/v2/rest/query")
 
     def execute(self, database: str, rendered_query: str, properties: ClientRequestProperties = None):
         self.executions.append((database, rendered_query, properties))
@@ -85,3 +88,43 @@ class TestTable(TestBase):
             [('test_db', 'test_table | take 5', None)],
             mock_kusto_client.executions,
         )
+
+    def test_cross_cluster_join(self):
+        mock_kusto_client_1 = MockKustoClient("https://one.kusto.windows.net")
+        mock_kusto_client_2 = MockKustoClient("https://two.kusto.windows.net")
+
+        table1 = PyKustoClient(mock_kusto_client_1)['test_db_1']['test_table_1']
+        table2 = PyKustoClient(mock_kusto_client_2)['test_db_2']['test_table_2']
+        Query(table1).take(5).join(Query(table2).take(6)).on(col.foo).execute()
+        self.assertEqual(
+            [('test_db_1', 'test_table_1 | take 5 | join  (cluster("two.kusto.windows.net").database("test_db_2").table("test_table_2") | take 6) on foo', None)],
+            mock_kusto_client_1.executions,
+        )
+
+    def test_cross_cluster_join_with_union(self):
+        mock_kusto_client_1 = MockKustoClient("https://one.kusto.windows.net")
+        mock_kusto_client_2 = MockKustoClient("https://two.kusto.windows.net")
+
+        table1 = PyKustoClient(mock_kusto_client_1)['test_db_1']['test_table_1']
+        table2 = PyKustoClient(mock_kusto_client_2)['test_db_2'].get_tables('test_table_2_*')
+        Query(table1).take(5).join(Query(table2).take(6)).on(col.foo).execute()
+        self.assertEqual(
+            [('test_db_1',
+              'test_table_1 | take 5 | join  (union cluster("two.kusto.windows.net").database("test_db_2").table("test_table_2_*") | take 6) on foo',
+              None)],
+            mock_kusto_client_1.executions,
+        )
+
+    def test_cross_cluster_join_with_union_2(self):
+        mock_kusto_client_1 = MockKustoClient("https://one.kusto.windows.net")
+        mock_kusto_client_2 = MockKustoClient("https://two.kusto.windows.net")
+
+        table1 = PyKustoClient(mock_kusto_client_1)['test_db_1']['test_table_1']
+        table2 = PyKustoClient(mock_kusto_client_2)['test_db_2'].get_tables('test_table_2_*', 'test_table_3_*')
+        Query(table1).take(5).join(Query(table2).take(6)).on(col.foo).execute()
+        self.assertEqual(
+            [('test_db_1',
+              'test_table_1 | take 5 | join  (union cluster("two.kusto.windows.net").database("test_db_2").table("test_table_2_*"), cluster("two.kusto.windows.net").database("test_db_2").table("test_table_3_*") | take 6) on foo',
+              None)],
+            mock_kusto_client_1.executions,
+        )