Views fk guesser

datafaker/db_utils.py

@@ -11,15 +11,7 @@
 
 # pylint: disable=no-name-in-module
 from psycopg2.errors import UndefinedObject  # ty: ignore[unresolved-import]
-from sqlalchemy import (
-    Column,
-    Connection,
-    Engine,
-    ForeignKey,
-    create_engine,
-    event,
-    select,
-)
+from sqlalchemy import Connection, Engine, ForeignKey, create_engine, event, select
 from sqlalchemy.engine.interfaces import DBAPIConnection
 from sqlalchemy.exc import (
     IntegrityError,
@@ -169,23 +161,32 @@ def create_db_engine(
         logger.error("DSN %s is malformed: %s", db_dsn, exc)
         raise Exit(1) from exc
 
-    settings = {}
+    settings: list[str] = []
+    dialect_name = get_sync_engine(engine).dialect.name
     if schema_name is not None:
-        if get_sync_engine(engine).dialect.name == "mssql":
-            engine = engine.execution_options(schema_translate_map={None: schema_name})
-        else:
-            settings["search_path"] = schema_name
+        match dialect_name:
+            case "mssql":
+                engine = engine.execution_options(
+                    schema_translate_map={None: schema_name}
+                )
+            case "snowflake":
+                settings.append(f"USE SCHEMA {schema_name};")
+            case _:
+                settings.append(f"SET search_path TO {schema_name};")
     dataset_views: list[str] = []
     if parquet_dir is not None:
-        tables = find_parquet_tables(parquet_dir)
-        joined = ",".join(
-            sorted({str(t.root.parent) for t in tables if not t.is_dataset})
-        )
-        # double up single quotes
-        dj = joined.replace("'", "''")
-        # enclose in single quotes
-        settings["file_search_path"] = f"'{dj}'"
-        dataset_views = [_dataset_view_sql(t) for t in tables if t.is_dataset]
+        match dialect_name:
+            case "duckdb":
+                tables = find_parquet_tables(parquet_dir)
+                # double up single quotes
+                dj = str(parquet_dir).replace("'", "''")
+                settings.append(f"SET file_search_path TO '{dj}';")
+                dataset_views = [_dataset_view_sql(t) for t in tables if t.hive_files()]
+            case _:
+                tables = []
+                logger.error(
+                    "--parquet-dir cannot be used with %s, only duckdb.", dialect_name
+                )
 
     if settings or dataset_views:
         event_engine = get_sync_engine(engine)
@@ -227,11 +228,15 @@ def create_db_engine_dst(
     return create_db_engine(db_dsn, schema_name, use_asyncio)
 
 
-def get_metadata(engine: Engine, schema_name: Optional[str] = None) -> MetaData:
+def get_metadata(
+    engine: Engine,
+    schema_name: Optional[str],
+    views_to_tables: bool,
+) -> MetaData:
     """Get the MetaData object associated with the engine passed."""
     md = MetaData()
     try:
-        md.reflect(engine, schema=schema_name)
+        md.reflect(engine, schema=schema_name, views=views_to_tables)
     except OperationalError as exc:
         logger.error("Cannot connect to database: %s", exc)
         raise Exit(1) from exc
@@ -245,7 +250,7 @@ def _sql_quote(text: str, quote: str) -> str:
 
 def _dataset_view_sql(table: ParquetTable) -> str:
     """Make the SQL that exposes a directory of parquet files as one table."""
-    files = ",".join(_sql_quote(str(f), "'") for f in table.files)
+    files = ",".join(_sql_quote(str(f), "'") for f in table.hive_files())
     return (
         f"CREATE OR REPLACE TEMP VIEW {_sql_quote(table.name, chr(34))} AS "
         f"SELECT * FROM read_parquet([{files}], "
@@ -267,16 +272,14 @@ def create_dataset_views(
     cursor.close()
 
 
-def set_db_settings(connection: DBAPIConnection, settings: Mapping[str, str]) -> None:
-    """Set the SEARCH_PATH for a PostgreSQL connection."""
+def set_db_settings(connection: DBAPIConnection, settings: list[str]) -> None:
+    """Run commands."""
     # https://docs.sqlalchemy.org/en/20/dialects/postgresql.html#remote-schema-table-introspection-and-postgresql-search-path
     existing_autocommit = connection.autocommit
     connection.autocommit = True
 
     cursor = connection.cursor()
-    # Parametrised queries don't work with asyncpg, hence the f-string.
-    sql = "".join(f"SET {k} TO {v};" for k, v in settings.items())
-    cursor.execute(sql)
+    cursor.execute("\n".join(settings))
     cursor.close()
 
     connection.autocommit = existing_autocommit
@@ -299,25 +302,6 @@ def get_orm_metadata(
     return metadata
 
 
-def get_fk_column_between(
-    from_: Table, to_: Table
-) -> tuple[Column, ForeignKey] | tuple[None, None]:
-    """
-    Get a foreign key between two tables.
-
-    :param from_: The table on which the foreign key should be defined.
-    :param to_: The table to which the foreign key should refer.
-    :return: A pair (column on from_, foreign key to to_) if such a key exists,
-     or (None, None) if not.
-    """
-    for c in from_.columns:
-        fk: ForeignKey
-        for fk in c.foreign_keys:
-            if fk.column.table == to_:
-                return (c, fk)
-    return (None, None)
-
-
 def is_numeric_sql_type(column_type: TypeEngine | None) -> bool:
     """Whether a SQLAlchemy column type is Numeric or Integer."""
     return isinstance(column_type, (Numeric, Integer))

datafaker/dialects.py

@@ -3,6 +3,7 @@
 from collections.abc import Mapping
 from typing import Any, Optional, TypeVar
 
+import snowflake.sqlalchemy.custom_types as snowflake_types
 from sqlalchemy import Column, Select, Table
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.schema import CreateSchema, CreateTable
@@ -515,3 +516,19 @@ def compile_random_mssql(_element: Random, _compiler: Any, **_kw: Any) -> str:
     same random number gets chosen for each row.
     """
     return "NEWID()"
+
+
+@compiles(snowflake_types.TIMESTAMP_LTZ)
+def compile_timestampltz_generic(
+    _element: snowflake_types.TIMESTAMP_LTZ, compiler: Any, **kw: Any
+) -> str:
+    """Render TIMESTAMP_LTZ as DateTime."""
+    return compiler.process(DateTime(timezone=False), **kw)
+
+
+@compiles(snowflake_types.TIMESTAMP_LTZ, "snowflake")
+def compile_timestampltz_ddb(
+    element: snowflake_types.TIMESTAMP_LTZ, compiler: Any, **kw: Any
+) -> str:
+    """Render TIMESTAMP_LTZ as TIMESTAMP_LTZ."""
+    return compiler.process(element, **kw)

datafaker/interactive/generators.py

@@ -8,7 +8,7 @@
 from enum import Enum
 from typing import Any, Callable, Optional, cast
 
-from sqlalchemy import Column, and_, literal_column, select
+from sqlalchemy import Column, and_, select
 from sqlalchemy.types import Integer
 
 from datafaker.db_utils import MaybeAsyncEngine, primary_private_fks, table_is_private
@@ -38,6 +38,7 @@
     schema_qualified_name,
     set_property,
     split_column_full_name,
+    text_list,
 )
 
 
@@ -161,7 +162,8 @@ def do(self, column_roles: RoleEntry | None) -> RoleEntry | None:
         """Perform the command."""
         if column_roles is None:
             return None
-        column_roles.new.remove(self.role)
+        if self.role in column_roles.new:
+            column_roles.new.remove(self.role)
         return column_roles
 
 
@@ -198,6 +200,9 @@ class GeneratorCmd(DbCmd):
     PRIMARY_PRIVATE_TEXT = "Primary Private"
     SECONDARY_PRIVATE_TEXT = "Secondary Private on columns {0}"
     NOT_PRIVATE_TEXT = "Not private"
+    REQUIRES_NO_SOURCE_DATA_TEXT = (
+        "{0}. {2}{1}{3} requires no data from the source database."
+    )
     REQUIRES_SOURCE_DATA_TEXT = (
         "{0}. {2}{1}{3} requires the following data from the source database:"
     )
@@ -472,7 +477,7 @@ def _copy_roles(self) -> None:
         """Set role information in the configuration."""
         for table_name, table_roles in self.roles.items():
             for column_name, entry in table_roles.items():
-                if entry.new:
+                if entry.new or entry.old:
                     set_property(
                         self.config,
                         ["tables", table_name, "columns", column_name, "roles"],
@@ -534,8 +539,8 @@ def _print_role_changes(self) -> int:
                         "Changing role set of column {0} of table {1} from {2} to {3}",
                         column_name,
                         table_name,
-                        ", ".join(str(oe) for oe in entry.old),
-                        ", ".join(str(ne) for ne in entry.new),
+                        text_list((e.value for e in entry.old), "empty"),
+                        text_list((e.value for e in entry.new), "empty"),
                     )
         return count
 
@@ -904,7 +909,7 @@ def _print_values_queried(self, table_name: str, n: int, prop: Proposer) -> None
         theme = get_active_theme()
         if not prop.select_aggregate_clauses() and not prop.custom_queries():
             self.print(
-                "{0}. {2}{1}{3} requires no data from the source database.",
+                self.REQUIRES_NO_SOURCE_DATA_TEXT,
                 n,
                 prop.name(),
                 theme.function,
@@ -1027,12 +1032,13 @@ def _print_select_aggregate_query(self, table_name: str, prop: Proposer) -> None
     def _get_column_data(
         self, count: int, to_str: Callable[[Any], str] = repr
     ) -> list[list[str]]:
-        columns = self._get_column_names()
-        col_exprs = [literal_column(col) for col in columns]
-        nonnull_clauses = [literal_column(col).isnot(None) for col in columns]
+        column_names = self._get_column_names()
+        table = self.table_metadata()
+        columns = [table.columns[c] for c in column_names]
+        nonnull_clauses = [col.isnot(None) for col in columns]
         stmt = (
-            select(*col_exprs)
-            .select_from(self.table_metadata())
+            select(*columns)
+            .select_from(table)
             .where(and_(*nonnull_clauses))
             .order_by(Random())
             .limit(count)

datafaker/interactive/missingness.py

@@ -215,7 +215,7 @@ def _copy_entries(self) -> None:
                         "kwargs": {
                             "patterns": f'SRC_STATS["{src_stat_key}"]["results"]'
                         },
-                        "columns": [c.name for c in entry.new_type.columns],
+                        "columns": [str(c.name) for c in entry.new_type.columns],
                     }
                 ]
                 src_stats.append(

datafaker/interactive/table.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from typing import Any, cast
 
-from sqlalchemy import func, literal_column, select
+from sqlalchemy import func, select
 
 from datafaker.dialects import Random
 from datafaker.interactive.base import (
@@ -469,26 +469,21 @@ def complete_data(
             return []
         return self.get_column_completions(text)
 
-    def print_column_data(self, column: str, count: int, min_length: int) -> None:
+    def print_column_data(self, column_name: str, count: int, min_length: int) -> None:
         """
         Print a sample of data from a certain column of the current table.
 
-        :param column: The name of the column to report on.
+        :param column_name: The name of the column to report on.
         :param count: The number of rows to sample.
         :param min_length: The minimum length of text to choose from (0 for any text).
         """
-        col_expr = literal_column(column)
+        table = self.table_metadata()
+        column = table.columns[column_name]
         if 0 < min_length:
-            where_clause = func.length(col_expr) >= min_length
+            where_clause = func.length(column) >= min_length
         else:
-            where_clause = col_expr.isnot(None)
-        stmt = (
-            select(col_expr)
-            .select_from(self.table_metadata())
-            .where(where_clause)
-            .order_by(Random())
-            .limit(count)
-        )
+            where_clause = column.isnot(None)
+        stmt = select(column).where(where_clause).order_by(Random()).limit(count)
         with self.sync_engine.connect() as connection:
             result = connection.execute(stmt)
             self.columnize([str(x[0]) for x in result.all()])

datafaker/main.py

@@ -172,11 +172,17 @@ def make_tables(
         help=(
             "Directory of Parquet files to consider part of the database."
             " This can be useful when using DuckDB."
+            " The .parquet files or directories containing partitioned tables"
+            " must all be directly inside this named directory."
             " Make sure you check the output!"
         ),
         file_okay=False,
         dir_okay=True,
     ),
+    views_to_tables: bool = Option(
+        False,
+        help=("Generate all views as tables in the synthetic version."),
+    ),
 ) -> None:
     """Make a YAML file representing the tables in the schema.
 
@@ -192,7 +198,9 @@ def make_tables(
     content = make_tables_file(
         get_source_dsn(),
         get_source_schema(),
-        parquet_dir,
+        views_to_tables=views_to_tables,
+        parquet_dir=parquet_dir,
+        guess_keys=views_to_tables or parquet_dir is not None,
     )
     orm_file_path.write_text(content, encoding="utf-8")
     logger.debug("%s created.", orm_file)