Skip to content

Commit

Permalink
Begin Table ACL migration
Browse files Browse the repository at this point in the history
Work in progress
  • Loading branch information
nfx committed Aug 15, 2023
1 parent b9886ef commit 85882ae
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 0 deletions.
12 changes: 12 additions & 0 deletions tests/integration/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import os
import pathlib
import random
import string
import sys
import uuid
from functools import partial
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand Down
57 changes: 57 additions & 0 deletions tests/integration/test_tacls.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 85882ae

Please sign in to comment.