fixup! cfr: Add support to export systables to a tarfile

cratedb_toolkit/cfr/cli.py

@@ -6,7 +6,7 @@
 import click
 from click_aliases import ClickAliasedGroup
 
-from cratedb_toolkit.cfr.systable import SystemTableExporter, SystemTableImporter
+from cratedb_toolkit.cfr.systable import SystemTableExporter, SystemTableImporter, Archive
 from cratedb_toolkit.util.cli import (
  boot_click,
  error_logger,
@@ -46,8 +46,20 @@ def cli(ctx: click.Context, cratedb_sqlalchemy_url: str, verbose: bool, debug: b
 def sys_export(ctx: click.Context, target: str):
  cratedb_sqlalchemy_url = ctx.meta["cratedb_sqlalchemy_url"]
  try:
- stc = SystemTableExporter(dburi=cratedb_sqlalchemy_url)
- path = stc.save(path_from_url(target))
+ target_path = path_from_url(target)
+ stc = SystemTableExporter(dburi=cratedb_sqlalchemy_url, target=target_path)
+
+ archive = None
+ if target_path.name.endswith(".tgz") or target_path.name.endswith(".tar.gz"):
+ archive = Archive(stc)
+
+ path = stc.save()
+
+ if archive is not None:
+ path = archive.make_tarfile()
+ archive.close()
+ logger.info(f"Created archive file {target}")
+
  jd({"path": str(path)})
  except Exception as ex:
  error_logger(ctx)(ex)

cratedb_toolkit/cfr/systable.py

@@ -90,12 +90,37 @@ def ddl(self, tablename_in: str, tablename_out: str, out_schema: str = None, wit
  return sql
 
 
-class SystemTableExporter:
+class PathProvider:
+
+ def __init__(self, path: t.Union[Path]):
+ self.path = path
+
+
+class Archive:
+
+ def __init__(self, path_provider: PathProvider):
+ self.path_provider = path_provider
+ self.temp_dir = tempfile.TemporaryDirectory()
+ self.target_path = self.path_provider.path
+ self.path_provider.path = Path(self.temp_dir.name)
+
+ def close(self):
+ self.temp_dir.cleanup()
+
+ def make_tarfile(self) -> Path:
+ source_path = self.path_provider.path
+ with tarfile.open(self.target_path, "x:gz") as tar:
+ tar.add(source_path.absolute(), arcname=os.path.basename(source_path))
+ return self.target_path
+
+
+class SystemTableExporter(PathProvider):
  """
  Export schema and data from CrateDB system tables.
  """
 
- def __init__(self, dburi: str, data_format: DataFormat = "jsonl"):
+ def __init__(self, dburi: str, target: t.Union[Path], data_format: DataFormat = "jsonl"):
+ super().__init__(target)
  self.dburi = dburi
  self.data_format = data_format
  self.adapter = DatabaseAdapter(dburi=self.dburi)
@@ -123,27 +148,10 @@ def dump_table(self, frame: pl.DataFrame, file: t.Union[t.TextIO, None] = None):
  else:
  raise NotImplementedError(f"Output format not implemented: {self.data_format}")
 
- def save(self, target: t.Union[Path]) -> Path:
- temp_dir = None
- if target.name.endswith(".tgz") or target.name.endswith(".tar.gz"):
- temp_dir = tempfile.TemporaryDirectory()
- target_folder = Path(temp_dir.name)
- else:
- target_folder = target
- target.mkdir(exist_ok=True, parents=True)
-
- full_path = self.export(target_folder)
-
- if temp_dir is not None:
- self.make_tarfile(target_folder, target)
- temp_dir.cleanup()
- logger.info(f"Created archive file {target}")
- return target
- return full_path
-
- def export(self, target_folder: Path) -> Path:
+ def save(self) -> Path:
+ self.path.mkdir(exist_ok=True, parents=True)
  timestamp = dt.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
- path = target_folder / self.info.cluster_name / timestamp / "sys"
+ path = self.path / self.info.cluster_name / timestamp / "sys"
  logger.info(f"Exporting system tables to: {path}")
  system_tables = self.inspector.table_names()
  path_schema = path / ExportSettings.SCHEMA_PATH
@@ -185,12 +193,6 @@ def export(self, target_folder: Path) -> Path:
  logger.info(f"Successfully exported {table_count} system tables")
  return path
 
- @staticmethod
- def make_tarfile(source_folder: Path, target_file_path: Path) -> Path:
- with tarfile.open(target_file_path, "x:gz") as tar:
- tar.add(source_folder.absolute(), arcname=os.path.basename(source_folder))
- return target_file_path
-
 
 class SystemTableImporter:
  """