This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch db-diagnostics in repository https://gitbox.apache.org/repos/asf/superset.git
commit 66214c6cd48695a2389862fc2e1fabe81c8497ca Author: Beto Dealmeida <[email protected]> AuthorDate: Tue Jul 25 16:55:55 2023 -0700 WIP --- superset/cli/test_db.py | 159 ++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 155 insertions(+), 4 deletions(-) diff --git a/superset/cli/test_db.py b/superset/cli/test_db.py index 30be55bcb9..d1442fc198 100644 --- a/superset/cli/test_db.py +++ b/superset/cli/test_db.py @@ -18,12 +18,25 @@ from __future__ import annotations import sys -from typing import Any, Dict, Type +from collections import defaultdict +from datetime import datetime +from typing import Any, Callable, Dict, Type import click import yaml from rich.console import Console -from sqlalchemy import create_engine +from sqlalchemy import ( + Column, + create_engine, + DateTime, + ForeignKey, + insert, + Integer, + MetaData, + select, + String, + Table, +) from sqlalchemy.engine import Engine from sqlalchemy.engine.url import make_url from sqlalchemy.exc import NoSuchModuleError @@ -80,6 +93,85 @@ ADVANCED_FEATURES = { "Supports validating SQL before running query": "sql_validation", } +metadata_obj = MetaData() + +user = Table( + "user", + metadata_obj, + Column("user_id", Integer, primary_key=True), + Column("user_name", String(16), nullable=False), + Column("email_address", String(60), key="email"), + Column("nickname", String(50), nullable=False), +) + +user_prefs = Table( + "user_prefs", + metadata_obj, + Column("pref_id", Integer, primary_key=True), + Column("user_id", Integer, ForeignKey("user.user_id"), nullable=False), + Column("pref_name", String(40), nullable=False), + Column("pref_value", String(100)), +) + + +TestType = Callable[[Console, Engine], None] + + +class TestRegistry: + def __init__(self) -> None: + self.tests: Dict[str, Any] = defaultdict(list) + + def add(self, *dialects: str) -> Callable[[TestType], TestType]: + def decorator(func: TestType) -> TestType: + for dialect in dialects: + self.tests[dialect].append(func) + + return func + + return decorator + + def get_tests(self, dialect: str) -> list[TestType]: + return self.tests[dialect] + + +registry = TestRegistry() + + [email protected]("sqlite", "postgresql") +def test_datetime(console: Console, engine: Engine) -> None: + """ + Create a table with a timestamp column. + """ + console.print("[bold]Testing datetime support...") + + md = MetaData() + table = Table( + "test", + md, + Column("ts", DateTime), + ) + + try: + console.print("Creating a table with a timestamp column...") + md.create_all(engine) + console.print("[green]Table created!") + + now = datetime.now() + + console.print("Inserting timestamp value...") + stmt = insert(table).values(ts=now) + engine.execute(stmt) + + console.print("Reading timestamp value...") + stmt = select(table) + row = engine.execute(stmt).fetchone() + assert row[0] == now + console.print(":thumbs_up: [green]Succcess!") + except Exception as ex: # pylint: disable=broad-except + console.print(f"[red]Test failed: {ex}") + console.print("[bold]Exiting...") + sys.exit(1) + @click.command() @click.argument("sqlalchemy_uri") @@ -300,5 +392,64 @@ def test_database_connectivity(console: Console, engine: Engine) -> None: color = "green" if result == 1 else "red" console.print(f"[{color}]> {result}") - # TODO(betodealmeida): create tables with different types and test them - # TODO(betodealmeida): run DB-specific tests + console.print("[bold]Checking that we can create tables...") + try: + metadata_obj.create_all(engine) + console.print("[green]Tables created!") + except Exception as ex: # pylint: disable=broad-except + console.print(f"[red]Unable to create tables: {ex}") + console.print("[bold]Exiting...") + sys.exit(1) + + console.print("[bold]Checking that we can insert data...") + stmt = insert(user).values( + user_name="beto", + email="[email protected]", + nickname="Beto", + ) + try: + console.print( + "sql>", + stmt.compile( + dialect=engine.dialect, + compile_kwargs={"literal_binds": True}, + ), + ) + engine.execute(stmt) + except Exception as ex: # pylint: disable=broad-except + console.print(f"[red]Unable to insert data: {ex}") + console.print("[bold]Exiting...") + sys.exit(1) + + console.print("[bold]Checking that we can read data...") + stmt = select(user).where(user.c.user_name == "beto") + try: + console.print( + "sql>", + stmt.compile( + dialect=engine.dialect, + compile_kwargs={"literal_binds": True}, + ), + ) + result = engine.execute(stmt).fetchall() + console.print(f"[green]> {result}") + except Exception as ex: # pylint: disable=broad-except + console.print(f"[red]Unable to read data: {ex}") + console.print("[bold]Exiting...") + sys.exit(1) + + console.print("[bold]Checking that we can drop tables...") + try: + metadata_obj.drop_all(engine) + console.print("[green]Done!") + except Exception as ex: # pylint: disable=broad-except + console.print(f"[red]Unable to drop tables: {ex}") + console.print("[bold]Exiting...") + sys.exit(1) + + # run engine-specific tests + tests = registry.get_tests(engine.dialect.name) + if tests: + console.print("[bold]Running engine-specific tests...") + for test in tests: + test(console, engine)
