cfr: Add ctk cfr diagnostics program

Add basic implementation for `sys-export` and `sys-import` subcommands.

It is about exporting system tables of CrateDB into SQL DDL and JSONL
files, and re-importing them for later analysis.

.gitignore

@@ -8,6 +8,7 @@ __pycache__
 dist
 .coverage*
 coverage.xml
+/cfr
 /foo
 /tmp
 /DOWNLOAD

CHANGES.md

@@ -2,7 +2,7 @@
 
 
 ## Unreleased
-- Add `cratedb-wtf` diagnostics program
+- Add `ctk cfr` and `ctk wtf` diagnostics programs
 
 ## 2024/04/10 v0.0.10
 - Dependencies: Unpin upper version bound of `dask`. Otherwise, compatibility

cratedb_toolkit/cfr/README.md

@@ -0,0 +1,28 @@
+# CrateDB Cluster Flight Recorder (CFR)
+
+Collect required cluster information for support requests
+and self-service debugging.
+
+
+## Synopsis
+
+Define CrateDB database cluster address.
+```shell
+export CRATEDB_SQLALCHEMY_URL=crate://localhost/
+```
+
+Export system table information into timestamped file,
+by default into the `cfr/sys` directory.
+```shell
+ctk cfr sys-export
+```
+
+Export system table information into given directory.
+```shell
+ctk cfr sys-export file:///var/ctk/cfr/sys
+```
+
+Import system table information from given directory.
+```shell
+ctk cfr sys-import file://./cfr/sys/2024-04-16T05-43-37
+```

cratedb_toolkit/cfr/backlog.md

@@ -0,0 +1,6 @@
+# CFR Backlog
+
+## Iteration +1
+- sys-export: Does the program need capabilities to **LIMIT** cardinality
+ on `sys-export` operations, for example, when they are super large?
+- sys-import: Accept target database schema. 

cratedb_toolkit/cfr/cli.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2021-2024, Crate.io Inc.
+# Distributed under the terms of the AGPLv3 license, see LICENSE.
+import logging
+import sys
+
+import click
+from click_aliases import ClickAliasedGroup
+
+from cratedb_toolkit.cfr.systable import SystemTableExporter, SystemTableImporter
+from cratedb_toolkit.util.cli import (
+ boot_click,
+ error_logger,
+ make_command,
+)
+from cratedb_toolkit.util.data import jd, path_from_url
+
+logger = logging.getLogger(__name__)
+
+
+cratedb_sqlalchemy_option = click.option(
+ "--cratedb-sqlalchemy-url", envvar="CRATEDB_SQLALCHEMY_URL", type=str, required=False, help="CrateDB SQLAlchemy URL"
+)
+
+
+@click.group(cls=ClickAliasedGroup) # type: ignore[arg-type]
+@cratedb_sqlalchemy_option
+@click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
+@click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
+@click.option("--scrub", envvar="SCRUB", is_flag=True, required=False, help="Blank out identifiable information")
+@click.version_option()
+@click.pass_context
+def cli(ctx: click.Context, cratedb_sqlalchemy_url: str, verbose: bool, debug: bool, scrub: bool):
+ """
+ Diagnostics and informational utilities.
+ """
+ if not cratedb_sqlalchemy_url:
+ logger.error("Unable to operate without database address")
+ sys.exit(1)
+ ctx.meta.update({"cratedb_sqlalchemy_url": cratedb_sqlalchemy_url, "scrub": scrub})
+ return boot_click(ctx, verbose, debug)
+
+
+@make_command(cli, "sys-export")
+@click.argument("target", envvar="CFR_TARGET", type=str, required=False, default="file://./cfr/sys")
+@click.pass_context
+def sys_export(ctx: click.Context, target: str):
+ cratedb_sqlalchemy_url = ctx.meta["cratedb_sqlalchemy_url"]
+ try:
+ stc = SystemTableExporter(dburi=cratedb_sqlalchemy_url, target=path_from_url(target))
+ path = stc.save()
+ jd({"path": str(path)})
+ except Exception as ex:
+ error_logger(ctx)(ex)
+ sys.exit(1)
+
+
+@make_command(cli, "sys-import")
+@click.argument("source", envvar="CFR_SOURCE", type=str, required=True)
+@click.pass_context
+def sys_import(ctx: click.Context, source: str):
+ cratedb_sqlalchemy_url = ctx.meta["cratedb_sqlalchemy_url"]
+ try:
+ stc = SystemTableImporter(dburi=cratedb_sqlalchemy_url, source=path_from_url(source))
+ stc.load()
+ except Exception as ex:
+ error_logger(ctx)(ex)
+ sys.exit(1)

cratedb_toolkit/cfr/systable.py

@@ -0,0 +1,210 @@
+"""
+CrateDB Diagnostics: System Tables Exporter and Importer.
+
+Schemas and results of following queries should be included:
+```sql
+SELECT * FROM sys.cluster
+SELECT * FROM sys.nodes
+SELECT * FROM sys.shards
+SELECT * FROM sys.allocations
+SELECT * FROM sys.jobs_log
+SELECT * FROM sys.operations_log
+```
+
+https://cratedb.com/docs/python/en/latest/by-example/sqlalchemy/inspection-reflection.html
+https://docs.sqlalchemy.org/en/20/faq/metadata_schema.html#how-can-i-get-the-create-table-drop-table-output-as-a-string
+"""
+
+import datetime as dt
+import logging
+import typing as t
+from pathlib import Path
+
+import polars as pl
+import sqlalchemy as sa
+from tqdm import tqdm
+
+from cratedb_toolkit.sqlalchemy.patch import patch_encoder
+from cratedb_toolkit.util import DatabaseAdapter
+from cratedb_toolkit.util.cli import error_logger
+
+logger = logging.getLogger(__name__)
+
+
+DataFormat = t.Literal["csv", "jsonl", "ndjson", "parquet"]
+
+
+class SystemTableKnowledge:
+ """
+ Manage a few bits of knowledge about CrateDB internals.
+ """
+
+ # Name of CrateDB's schema for system tables.
+ SYS_SCHEMA = "sys"
+
+ # TODO: Reflecting the `summits` table raises an error.
+ # AttributeError: 'UserDefinedType' object has no attribute 'get_col_spec'
+ REFLECTION_BLOCKLIST = ["summits"]
+
+
+class ExportSettings:
+ """
+ Manage a few bits of knowledge about how to export system tables from CrateDB.
+ """
+
+ # Subdirectories where to store schema vs. data information.
+ SCHEMA_PATH = "schema"
+ DATA_PATH = "data"
+
+ # The filename prefix when storing tables to disk.
+ TABLE_FILENAME_PREFIX = "sys-"
+
+
+class SystemTableInspector:
+ """
+ Reflect schema information from CrateDB system tables.
+ """
+
+ def __init__(self, dburi: str):
+ self.dburi = dburi
+ self.adapter = DatabaseAdapter(dburi=self.dburi)
+ self.engine = self.adapter.engine
+ self.inspector = sa.inspect(self.engine)
+
+ def table_names(self):
+ return self.inspector.get_table_names(schema=SystemTableKnowledge.SYS_SCHEMA)
+
+ def ddl(self, tablename_in: str, tablename_out: str, out_schema: str = None, with_drop_table: bool = False) -> str:
+ meta = sa.MetaData(schema=SystemTableKnowledge.SYS_SCHEMA)
+ table = sa.Table(tablename_in, meta, autoload_with=self.engine)
+ table.schema = out_schema
+ table.name = tablename_out
+ sql = ""
+ if with_drop_table:
+ sql += sa.schema.DropTable(table, if_exists=True).compile(self.engine).string.strip() + ";\n"
+ sql += sa.schema.CreateTable(table, if_not_exists=True).compile(self.engine).string.strip() + ";\n"
+ return sql
+
+
+class SystemTableExporter:
+ """
+ Export schema and data from CrateDB system tables.
+ """
+
+ def __init__(self, dburi: str, target: t.Union[Path], data_format: DataFormat = "jsonl"):
+ self.dburi = dburi
+ self.target = target
+ self.data_format = data_format
+ self.adapter = DatabaseAdapter(dburi=self.dburi)
+ self.engine = self.adapter.engine
+ self.inspector = SystemTableInspector(dburi=self.dburi)
+ self.target.mkdir(exist_ok=True, parents=True)
+
+ def dump_table(self, tablename: str, file: t.Union[t.TextIO, None] = None):
+ sql = f'SELECT * FROM "{SystemTableKnowledge.SYS_SCHEMA}"."{tablename}"' # noqa: S608
+ # logger.info(f"Running SQL: {sql}") # noqa: ERA001
+ df: pl.DataFrame = pl.read_database(
+ query=sql, # noqa: S608
+ connection=self.engine,
+ )
+ if self.data_format == "csv":
+ # polars.exceptions.ComputeError: CSV format does not support nested data
+ # return df.write_csv() # noqa: ERA001
+ return df.to_pandas().to_csv(file)
+ elif self.data_format in ["jsonl", "ndjson"]:
+ return df.write_ndjson(file and file.buffer) # type: ignore[arg-type]
+ elif self.data_format in ["parquet", "pq"]:
+ return df.write_parquet(file and file.buffer) # type: ignore[arg-type]
+ else:
+ raise NotImplementedError(f"Output format not implemented: {self.data_format}")
+
+ def save(self) -> Path:
+ timestamp = dt.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+ path = self.target / timestamp
+ logger.info(f"Exporting system tables to: {path}")
+ system_tables = self.inspector.table_names()
+ path_schema = path / ExportSettings.SCHEMA_PATH
+ path_data = path / ExportSettings.DATA_PATH
+ path_schema.mkdir(parents=True, exist_ok=True)
+ path_data.mkdir(parents=True, exist_ok=True)
+ table_count = 0
+ for tablename in tqdm(system_tables, disable=None):
+ if tablename in SystemTableKnowledge.REFLECTION_BLOCKLIST:
+ continue
+ path_table_schema = path_schema / f"{ExportSettings.TABLE_FILENAME_PREFIX}{tablename}.sql"
+ path_table_data = path_data / f"{ExportSettings.TABLE_FILENAME_PREFIX}{tablename}.{self.data_format}"
+ tablename_out = f"{ExportSettings.TABLE_FILENAME_PREFIX}{tablename}"
+ with open(path_table_schema, "w") as fh_schema:
+ print(self.inspector.ddl(tablename_in=tablename, tablename_out=tablename_out), file=fh_schema)
+ mode = "w"
+ if self.data_format in ["parquet", "pq"]:
+ mode = "wb"
+ with open(path_table_data, mode) as fh_data:
+ self.dump_table(tablename=tablename, file=t.cast(t.TextIO, fh_data))
+ table_count += 1
+ logger.info(f"Successfully exported {table_count} system tables")
+ return path
+
+
+class SystemTableImporter:
+ """
+ Import schema and data about CrateDB system tables.
+ """
+
+ def __init__(self, dburi: str, source: Path, data_format: DataFormat = "jsonl", debug: bool = False):
+ self.dburi = dburi
+ self.source = source
+ self.data_format = data_format
+ self.debug = debug
+ self.adapter = DatabaseAdapter(dburi=self.dburi)
+
+ def table_names(self):
+ path_schema = self.source / ExportSettings.SCHEMA_PATH
+ names = []
+ for item in path_schema.glob("*.sql"):
+ name = item.name.replace(ExportSettings.TABLE_FILENAME_PREFIX, "").replace(".sql", "")
+ names.append(name)
+ return names
+
+ def load(self):
+ path_schema = self.source / ExportSettings.SCHEMA_PATH
+ path_data = self.source / ExportSettings.DATA_PATH
+
+ if not path_schema.exists():
+ raise FileNotFoundError(f"Path does not exist: {path_schema}")
+
+ logger.info(f"Importing system tables from: {self.source}")
+
+ for tablename in tqdm(self.table_names()):
+ tablename_restored = ExportSettings.TABLE_FILENAME_PREFIX + tablename
+
+ path_table_schema = path_schema / f"{ExportSettings.TABLE_FILENAME_PREFIX}{tablename}.sql"
+ path_table_data = path_data / f"{ExportSettings.TABLE_FILENAME_PREFIX}{tablename}.{self.data_format}"
+
+ # Skip empty files.
+ if Path(path_table_data).stat().st_size == 0:
+ continue
+
+ # Invoke SQL DDL.
+ schema_sql = path_table_schema.read_text()
+ self.adapter.run_sql(schema_sql)
+
+ # Load data.
+ try:
+ df: pl.DataFrame = self.load_table(path_table_data)
+ df.write_database(table_name=tablename_restored, connection=self.dburi, if_table_exists="append")
+ except Exception as ex:
+ error_logger(self.debug)(f"Importing table failed: {tablename}. Reason: {ex}")
+
+ # df.to_pandas().to_sql(name=tablename, con=self.adapter.engine, if_exists="append", index=False) # noqa: ERA001, E501
+
+ def load_table(self, path: Path) -> pl.DataFrame:
+ if path.suffix in [".jsonl"]:
+ return pl.read_ndjson(path)
+ elif path.suffix in [".parquet", ".pq"]:
+ return pl.read_parquet(path)
+ else:
+ raise NotImplementedError(f"Input format not implemented: {path.suffix}")
+
+
+patch_encoder()

cratedb_toolkit/cli.py

@@ -3,6 +3,7 @@
 
 from cratedb_toolkit.util.cli import boot_click
 
+from .cfr.cli import cli as cfr_cli
 from .cluster.cli import cli as cloud_cli
 from .io.cli import cli as io_cli
 from .job.cli import cli_list_jobs
@@ -19,6 +20,7 @@ def cli(ctx: click.Context, verbose: bool, debug: bool):
  return boot_click(ctx, verbose, debug)
 
 
+cli.add_command(cfr_cli, name="cfr")
 cli.add_command(cloud_cli, name="cluster")
 cli.add_command(io_cli, name="load")
 cli.add_command(shell_cli, name="shell")

cratedb_toolkit/sqlalchemy/patch.py

@@ -1,5 +1,11 @@
+import calendar
+import datetime as dt
+import json
 import typing as t
+from decimal import Decimal
+from uuid import UUID
 
+import numpy as np
 import sqlalchemy as sa
 
 
@@ -32,3 +38,37 @@ def get_table_names(self, connection: sa.Connection, schema: t.Optional[str] = N
  return get_table_names_dist(self, connection=connection, schema=schema, **kw)
 
  CrateDialect.get_table_names = get_table_names # type: ignore
+
+
+def patch_encoder():
+ import crate.client.http
+
+ crate.client.http.CrateJsonEncoder = CrateJsonEncoderWithNumPy
+
+
+class CrateJsonEncoderWithNumPy(json.JSONEncoder):
+ epoch_aware = dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc)
+ epoch_naive = dt.datetime(1970, 1, 1)
+
+ def default(self, o):
+ # Vanilla CrateDB Python.
+ if isinstance(o, (Decimal, UUID)):
+ return str(o)
+ if isinstance(o, dt.datetime):
+ if o.tzinfo is not None:
+ delta = o - self.epoch_aware
+ else:
+ delta = o - self.epoch_naive
+ return int(delta.microseconds / 1000.0 + (delta.seconds + delta.days * 24 * 3600) * 1000.0)
+ if isinstance(o, dt.date):
+ return calendar.timegm(o.timetuple()) * 1000
+
+ # NumPy ndarray and friends.
+ # https://stackoverflow.com/a/49677241
+ if isinstance(o, np.integer):
+ return int(o)
+ elif isinstance(o, np.floating):
+ return float(o)
+ elif isinstance(o, np.ndarray):
+ return o.tolist()
+ return json.JSONEncoder.default(self, o)