From 85882ae40910ff5eefd0d90efe515245a16bc428 Mon Sep 17 00:00:00 2001
From: Serge Smertin <serge.smertin@databricks.com>
Date: Thu, 10 Aug 2023 21:24:35 +0200
Subject: [PATCH] Begin Table ACL migration

Work in progress
---
 tests/integration/conftest.py   | 12 +++++++
 tests/integration/test_tacls.py | 57 +++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)
 create mode 100644 tests/integration/test_tacls.py

diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index c53843b72c..276fb4e041 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -4,6 +4,7 @@
 import os
 import pathlib
 import random
+import string
 import sys
 import uuid
 from functools import partial
@@ -108,6 +109,16 @@ def _is_in_debug() -> bool:
     ]
 
 
+@pytest.fixture
+def random():
+    import random
+
+    def inner(k=16) -> str:
+        charset = string.ascii_uppercase + string.ascii_lowercase + string.digits
+        return ''.join(random.choices(charset, k=int(k)))
+
+    return inner
+
 @pytest.fixture(scope="session")
 def ws() -> ImprovedWorkspaceClient:
     # Use variables from Unified Auth
@@ -134,6 +145,7 @@ def account_host(cfg: Config) -> str:
 
 @pytest.fixture(scope="session")
 def dbconnect_cluster_id(ws: ImprovedWorkspaceClient) -> str:
+    # TODO: will use predeclared DATABRICKS_CLUSTER_ID env variable
     dbc_cluster = next(filter(lambda c: c.cluster_name == DB_CONNECT_CLUSTER_NAME, ws.clusters.list()), None)
 
     if dbc_cluster:
diff --git a/tests/integration/test_tacls.py b/tests/integration/test_tacls.py
new file mode 100644
index 0000000000..1db1a01be5
--- /dev/null
+++ b/tests/integration/test_tacls.py
@@ -0,0 +1,57 @@
+import os
+from dataclasses import dataclass
+from typing import Type
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service import compute
+
+from uc_migration_toolkit.providers.client import ImprovedWorkspaceClient
+from uc_migration_toolkit.providers.logger import logger
+
+
+@dataclass
+class ShowTable:
+    database: str
+    tableName: str
+    isTemporary: bool
+
+
+class CommandExecutor:
+    def __init__(self, ws: WorkspaceClient, cluster_id):
+        self.cluster_id = cluster_id
+        self.ws = ws
+        self.ctx = ws.command_execution.create(
+            cluster_id=cluster_id,
+            language=compute.Language.PYTHON).result()
+
+    def run(self, command: str, language=compute.Language.PYTHON):
+        return self.ws.command_execution.execute(
+            cluster_id=self.cluster_id,
+            context_id=self.ctx.id,
+            language=language,
+            command=command,
+        ).result()
+
+    def sql(self, query: str, result_type: Type = None):
+        logger.debug(f'Running SQL query on {self.cluster_id}: {query}')
+        res = self.run(query, language=compute.Language.SQL)
+        results = res.results
+        if result_type is None:
+            schema_str = ', '.join([f'{c["name"]}:{c["type"]}' for c in results.schema])
+            raise ValueError(f'Needs constructor with schema: {schema_str}')
+        if results.result_type == compute.ResultType.TABLE:
+            for tpl in results.data:
+                yield result_type(*tpl)
+        else:
+            raise RuntimeError(f'Unknown result type: {results.result_type}: {results.summary}')
+
+
+def test_table_acls(ws: ImprovedWorkspaceClient, random):
+    cluster_id = os.environ['TABLE_ACL_CLUSTER_ID']
+    ce = CommandExecutor(ws, cluster_id)
+
+    table_name = f'ucx_{random(12)}'
+    ce.sql(f'CREATE TABLE {table_name} AS SELECT 2+2 AS four, 2+3 AS five')
+
+    for t in ce.sql("SHOW TABLES", ShowTable):
+        print(t)