mirror of
https://github.com/hanzo-ds/connect.git
synced 2026-10-30 20:31:24 +00:00
bec10356fe
* refactor(cc_sqlalchemy): add type annotations * refactor(cc_sqlalchemy): make TableEngine inherit SchemaItem directly SchemaItem is SchemaEventTarget + Visitable formalized — TableEngine was already structurally identical but didn't declare the inheritance, causing a # type: ignore[arg-type] suppression wherever a TableEngine instance was passed to Table(). Switching to SchemaItem makes the relationship explicit and removes the suppression in alembic/impl.py. Visitable.__init__ was a no-op (delegates to object.__init__); replaced with super().__init__(). * formatting --------- Co-authored-by: Joe S <[email protected]>
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
from sqlalchemy.exc import ArgumentError
|
|
from sqlalchemy.sql.ddl import DDL
|
|
|
|
from clickhouse_connect.driver.binding import format_str, quote_identifier
|
|
|
|
|
|
class CreateDatabase(DDL):
|
|
"""
|
|
SqlAlchemy DDL statement that is essentially an alternative to the built in CreateSchema DDL class
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
engine: str | None = None,
|
|
zoo_path: str | None = None,
|
|
shard_name: str = "{shard}",
|
|
replica_name: str = "{replica}",
|
|
exists_ok: bool = False,
|
|
):
|
|
"""
|
|
:param name: Database name
|
|
:param engine: Database ClickHouse engine type
|
|
:param zoo_path: ClickHouse zookeeper path for Replicated database engine
|
|
:param shard_name: Clickhouse shard name for Replicated database engine
|
|
:param replica_name: Replica name for Replicated database engine
|
|
"""
|
|
if engine and engine not in ("Ordinary", "Atomic", "Lazy", "Replicated"):
|
|
raise ArgumentError(f"Unrecognized engine type {engine}")
|
|
stmt = f"CREATE DATABASE {'IF NOT EXISTS ' if exists_ok else ''}{quote_identifier(name)}"
|
|
if engine:
|
|
stmt += f" Engine {engine}"
|
|
if engine == "Replicated":
|
|
if not zoo_path:
|
|
raise ArgumentError("zoo_path is required for Replicated Database Engine")
|
|
stmt += f" ({format_str(zoo_path)}, {format_str(shard_name)}, {format_str(replica_name)})"
|
|
super().__init__(stmt)
|
|
|
|
|
|
class DropDatabase(DDL):
|
|
"""
|
|
Alternative DDL statement for built in SqlAlchemy DropSchema DDL class
|
|
"""
|
|
|
|
def __init__(self, name: str, missing_ok: bool = False):
|
|
super().__init__(f"DROP DATABASE {'IF EXISTS ' if missing_ok else ''}{quote_identifier(name)}")
|