Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 14 additions & 12 deletions datafaker/interactive/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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``.
Expand All @@ -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]:
"""
Expand All @@ -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
Expand Down
41 changes: 31 additions & 10 deletions datafaker/interactive/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down
9 changes: 3 additions & 6 deletions datafaker/interactive/generators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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``.
Expand All @@ -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
Expand Down
16 changes: 4 additions & 12 deletions datafaker/interactive/missingness.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down Expand Up @@ -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
Expand Down
10 changes: 3 additions & 7 deletions datafaker/interactive/table.py
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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
Expand Down
Loading