Spaces:
Runtime error
Runtime error
| from typing import Sequence, Optional, Union | |
| from clickhouse_connect.driver import Client | |
| class TableContext: | |
| def __init__(self, client: Client, | |
| table: str, | |
| columns: Union[str, Sequence[str]], | |
| column_types: Optional[Sequence[str]] = None, | |
| engine: str = 'MergeTree', | |
| order_by: str = None): | |
| self.client = client | |
| self.table = table | |
| if isinstance(columns, str): | |
| columns = columns.split(',') | |
| if column_types is None: | |
| self.column_names = [] | |
| self.column_types = [] | |
| for col in columns: | |
| col = col.strip() | |
| ix = col.find(' ') | |
| self.column_types.append(col[ix + 1:].strip()) | |
| self.column_names.append(col[:ix].strip()) | |
| else: | |
| self.column_names = columns | |
| self.column_types = column_types | |
| self.engine = engine | |
| self.order_by = self.column_names[0] if order_by is None else order_by | |
| def __enter__(self): | |
| if self.client.min_version('19'): | |
| self.client.command(f'DROP TABLE IF EXISTS {self.table}') | |
| else: | |
| self.client.command(f'DROP TABLE IF EXISTS {self.table} SYNC') | |
| col_defs = ','.join(f'{name} {col_type}' for name, col_type in zip(self.column_names, self.column_types)) | |
| self.client.command(f'CREATE TABLE {self.table} ({col_defs}) ENGINE {self.engine} ORDER BY {self.order_by}') | |
| return self | |
| def __exit__(self, exc_type, exc_val, exc_tb): | |
| self.client.command(f'DROP TABLE IF EXISTS {self.table}') | |