ClickHouse/tests/testflows/rbac/helper/common.py

import uuid

import testflows.settings as settings

from contextlib import contextmanager

from testflows.core.name import basename, parentname
from testflows._core.testtype import TestSubType
from testflows.core import *

from helpers.common import instrument_clickhouse_server_log
from rbac.helper.tables import table_types


def permutations(table_count=1):
    return [*range((1 << table_count) - 1)]


def getuid():
    if current().subtype == TestSubType.Example:
        testname = (
            f"{basename(parentname(current().name)).replace(' ', '_').replace(',','')}"
        )
    else:
        testname = f"{basename(current().name).replace(' ', '_').replace(',','')}"
    return testname + "_" + str(uuid.uuid1()).replace("-", "_")


@contextmanager
def table(node, name, table_type_name="MergeTree"):
    table_type = table_types[table_type_name]
    try:
        names = name.split(",")
        for name in names:
            with Given(f"I have {name} with engine {table_type_name}"):
                node.query(f"DROP TABLE IF EXISTS {name}")
                node.query(table_type.create_statement.format(name=name))
        yield
    finally:
        for name in names:
            with Finally(f"I drop the table {name}"):
                if table_type.cluster:
                    node.query(
                        f"DROP TABLE IF EXISTS {name} ON CLUSTER {table_type.cluster}"
                    )
                else:
                    node.query(f"DROP TABLE IF EXISTS {name}")


@contextmanager
def user(node, name):
    try:
        names = name.split(",")
        for name in names:
            with Given("I have a user"):
                node.query(f"CREATE USER OR REPLACE {name}")
        yield
    finally:
        for name in names:
            with Finally("I drop the user"):
                node.query(f"DROP USER IF EXISTS {name}")


@contextmanager
def role(node, role):
    try:
        roles = role.split(",")
        for role in roles:
            with Given("I have a role"):
                node.query(f"CREATE ROLE OR REPLACE {role}")
        yield
    finally:
        for role in roles:
            with Finally("I drop the role"):
                node.query(f"DROP ROLE IF EXISTS {role}")


@TestStep(Given)
def row_policy(self, name, table, node=None):
    """Create a row policy with a given name on a given table."""
    if node is None:
        node = self.context.node

    try:
        with Given(f"I create row policy {name}"):
            node.query(f"CREATE ROW POLICY {name} ON {table}")
        yield

    finally:
        with Finally(f"I delete row policy {name}"):
            node.query(f"DROP ROW POLICY IF EXISTS {name} ON {table}")


tables = {
    "table0": 1 << 0,
    "table1": 1 << 1,
    "table2": 1 << 2,
    "table3": 1 << 3,
    "table4": 1 << 4,
    "table5": 1 << 5,
    "table6": 1 << 6,
    "table7": 1 << 7,
}


@contextmanager
def grant_select_on_table(node, grants, target_name, *table_names):
    try:
        tables_granted = []
        for table_number in range(len(table_names)):

            if grants & tables[f"table{table_number}"]:

                with When(f"I grant select privilege on {table_names[table_number]}"):
                    node.query(
                        f"GRANT SELECT ON {table_names[table_number]} TO {target_name}"
                    )

                    tables_granted.append(f"{table_names[table_number]}")

        yield (", ").join(tables_granted)

    finally:
        for table_number in range(len(table_names)):
            with Finally(
                f"I revoke the select privilege on {table_names[table_number]}"
            ):
                node.query(
                    f"REVOKE SELECT ON {table_names[table_number]} FROM {target_name}"
                )