Parquet dir stored so it can be used by all source commands. (#83)

#82 make-stats gains a --parquet-dir option
refactoring so everything can work with parquet

Author: tim-band

datafaker/interactive/__init__.py

@@ -7,6 +7,7 @@
 
 from sqlalchemy import MetaData
 
+from datafaker.interactive.base import DbCmd
 from datafaker.interactive.generators import GeneratorCmd, try_setting_generator
 from datafaker.interactive.missingness import MissingnessCmd
 from datafaker.interactive.table import TableCmd
@@ -25,10 +26,15 @@
 
 
 def update_config_tables(
-    src_dsn: str, src_schema: str | None, metadata: MetaData, config: MutableMapping
+    src_dsn: str,
+    src_schema: str | None,
+    metadata: MetaData,
+    config: MutableMapping,
+    parquet_dir: Path | None,
 ) -> Mapping[str, Any]:
     """Ask the user to specify what should happen to each table."""
-    with TableCmd(src_dsn, src_schema, metadata, config) as tc:
+    settings = DbCmd.Settings(src_dsn, src_schema, config, metadata, parquet_dir)
+    with TableCmd(settings) as tc:
         tc.cmdloop()
         return tc.config
 
@@ -38,6 +44,7 @@ def update_missingness(
     src_schema: str | None,
     metadata: MetaData,
     config: MutableMapping[str, Any],
+    parquet_dir: Path | None,
 ) -> Mapping[str, Any]:
     """
     Ask the user to update the missingness information in ``config.yaml``.
@@ -49,16 +56,14 @@ def update_missingness(
     :param config: The starting configuration,
     :return: The updated configuration.
     """
-    with MissingnessCmd(src_dsn, src_schema, metadata, config) as mc:
+    settings = DbCmd.Settings(src_dsn, src_schema, config, metadata, parquet_dir)
+    with MissingnessCmd(settings) as mc:
         mc.cmdloop()
         return mc.config
 
 
 def update_config_generators(
-    src_dsn: str,
-    src_schema: str | None,
-    metadata: MetaData,
-    config: MutableMapping[str, Any],
+    settings: DbCmd.Settings,
     spec_path: Path | None,
 ) -> Mapping[str, Any]:
     """
@@ -68,14 +73,11 @@ def update_config_generators(
     Column name (or space-separated list of column names), Generator
     name required, Second choice generator name, Third choice generator
     name, etcetera.
-    :param src_dsn: Address of the source database
-    :param src_schema: Name of the source database schema to read from
-    :param metadata: SQLAlchemy representation of the source database
-    :param config: Existing configuration (will be destructively updated)
+    :param settings: Source database settings.
     :param spec_path: The path of the CSV file containing the specification
     :return: Updated configuration.
     """
-    with GeneratorCmd(src_dsn, src_schema, metadata, config) as gc:
+    with GeneratorCmd(settings) as gc:
         if spec_path is None:
             gc.cmdloop()
             return gc.config

datafaker/interactive/base.py

@@ -4,6 +4,7 @@
 from collections.abc import Mapping, MutableMapping, Sequence
 from dataclasses import dataclass
 from enum import Enum
+from pathlib import Path
 from types import TracebackType
 from typing import Any, Optional, Type
 
@@ -121,30 +122,50 @@ def make_table_entry(
         :return: The table entry or None if this table should not be interacted with.
         """
 
+    @dataclass
+    class Settings:
+        """Settings for the source database."""
+
+        dsn: str
+        schema: str | None
+        config: MutableMapping[str, Any]
+        metadata: MetaData
+        parquet_dir: Path | None
+
     def __init__(
         self,
-        src_dsn: str,
-        src_schema: str | None,
-        metadata: MetaData,
-        config: MutableMapping[str, Any],
+        settings: Settings,
     ):
-        """Initialise a DbCmd."""
+        """
+        Initialise a DbCmd.
+
+        :param src_dsn: The database connection string for the source database.
+        :param src_schema: The name of the schema name for the source database.
+        :param metadata: The metadata for the source database.
+        :param config: The ``config.xml`` object.
+        :param parquet_dir: The directory where parquet files are stored that
+          are to be considered part of the source database (only for DuckDB).
+        """
         super().__init__()
-        self.config: MutableMapping[str, Any] = config
-        self.metadata = metadata
+        self.config: MutableMapping[str, Any] = settings.config
+        self.metadata = settings.metadata
         self._table_entries: list[TableEntry] = []
-        tables_config: MutableMapping = config.get("tables", {})
+        tables_config: MutableMapping = self.config.get("tables", {})
         if not isinstance(tables_config, MutableMapping):
             tables_config = {}
-        for name in metadata.tables.keys():
+        for name in self.metadata.tables.keys():
             table_config = tables_config.get(name, {})
             if not isinstance(table_config, MutableMapping):
                 table_config = {}
             entry = self.make_table_entry(name, table_config)
             if entry is not None:
                 self._table_entries.append(entry)
         self.table_index = 0
-        self.engine = create_db_engine(src_dsn, schema_name=src_schema)
+        self.engine = create_db_engine(
+            settings.dsn,
+            schema_name=settings.schema,
+            parquet_dir=settings.parquet_dir,
+        )
 
     @property
     def sync_engine(self) -> Engine:

datafaker/interactive/generators.py

@@ -6,7 +6,7 @@
 from typing import Any, Callable, Optional, cast
 
 import sqlalchemy
-from sqlalchemy import Column, MetaData
+from sqlalchemy import Column
 
 from datafaker.generators import everything_factory
 from datafaker.generators.base import Generator, PredefinedGenerator
@@ -147,10 +147,7 @@ def make_table_entry(
 
     def __init__(
         self,
-        src_dsn: str,
-        src_schema: str | None,
-        metadata: MetaData,
-        config: MutableMapping[str, Any],
+        settings: DbCmd.Settings,
     ) -> None:
         """
         Initialise a ``GeneratorCmd``.
@@ -160,7 +157,7 @@ def __init__(
         :param metadata: SQLAlchemy metadata for the source database
         :param config: Configuration loaded from ``config.yaml``
         """
-        super().__init__(src_dsn, src_schema, metadata, config)
+        super().__init__(settings)
         self.generators: list[Generator] | None = None
         self.generator_index = 0
         self.generators_valid_columns: Optional[tuple[int, list[str]]] = None

datafaker/interactive/missingness.py

@@ -1,11 +1,9 @@
 """Missingness configuration shell."""
 import re
-from collections.abc import Iterable, Mapping, MutableMapping
+from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
 from typing import cast
 
-from sqlalchemy import MetaData
-
 from datafaker.interactive.base import DbCmd, TableEntry
 
 
@@ -139,20 +137,14 @@ def make_table_entry(
 
     def __init__(
         self,
-        src_dsn: str,
-        src_schema: str | None,
-        metadata: MetaData,
-        config: MutableMapping,
+        settings: DbCmd.Settings,
     ):
         """
         Initialise a MissingnessCmd.
 
-        :param src_dsn: connection string for the source database.
-        :param src_schema: schema name for the source database.
-        :param metadata: SQLAlchemy metadata for the source database.
-        :param config: Configuration from the ``config.yaml`` file.
+        :param settings: source database settings.
         """
-        super().__init__(src_dsn, src_schema, metadata, config)
+        super().__init__(settings)
         self.set_prompt()
 
     @property

datafaker/interactive/table.py

@@ -1,10 +1,9 @@
 """Table configuration command shell."""
-from collections.abc import Mapping, MutableMapping
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any, cast
 
 import sqlalchemy
-from sqlalchemy import MetaData
 
 from datafaker.interactive.base import (
     TYPE_LETTER,
@@ -92,15 +91,12 @@ def make_table_entry(
 
     def __init__(
         self,
-        src_dsn: str,
-        src_schema: str | None,
-        metadata: MetaData,
-        config: MutableMapping[str, Any],
+        settings: DbCmd.Settings,
         *args: Any,
         **kwargs: Any,
     ) -> None:
         """Initialise a TableCmd."""
-        super().__init__(src_dsn, src_schema, metadata, config, *args, **kwargs)
+        super().__init__(settings, *args, **kwargs)
         self.set_prompt()
 
     @property