mirror of
https://github.com/PostHog/posthog.git
synced 2024-11-30 19:41:46 +01:00
30 lines
1.2 KiB
Python
30 lines
1.2 KiB
Python
import logging
|
|
|
|
from typing import Any
|
|
from collections.abc import Callable
|
|
from django.core.management.base import BaseCommand, CommandParser
|
|
|
|
from posthog.clickhouse.materialized_columns import ColumnName, TablesWithMaterializedColumns
|
|
from ee.clickhouse.materialized_columns.columns import update_column_is_disabled, drop_column
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
COLUMN_OPERATIONS: dict[str, Callable[[TablesWithMaterializedColumns, ColumnName], Any]] = {
|
|
"enable": lambda table, column_name: update_column_is_disabled(table, column_name, is_disabled=False),
|
|
"disable": lambda table, column_name: update_column_is_disabled(table, column_name, is_disabled=True),
|
|
"drop": drop_column,
|
|
}
|
|
|
|
|
|
class Command(BaseCommand):
|
|
def add_arguments(self, parser: CommandParser) -> None:
|
|
parser.add_argument("operation", choices=COLUMN_OPERATIONS.keys())
|
|
parser.add_argument("table")
|
|
parser.add_argument("column_name")
|
|
|
|
def handle(self, operation: str, table: TablesWithMaterializedColumns, column_name: ColumnName, **options):
|
|
logger.info("Running %r for %r.%r...", operation, table, column_name)
|
|
fn = COLUMN_OPERATIONS[operation]
|
|
fn(table, column_name)
|
|
logger.info("Success!")
|