Initial Snowflake support

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,22 +161,34 @@ 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}'"
+        match dialect_name:
+            case "duckdb":
+                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("'", "''")
+                settings.append(f"SET file_search_path TO '{dj}';")
+            case _:
+                tables = []
+                logger.error(
+                    "--parquet-dir cannot be used with %s, only duckdb.", dialect_name
+                )
         dataset_views = [_dataset_view_sql(t) for t in tables if t.is_dataset]
 
     if settings or dataset_views:
@@ -227,11 +231,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
@@ -267,16 +275,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 +305,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

@@ -177,6 +177,10 @@ def make_tables(
         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 +196,8 @@ 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,
     )
     orm_file_path.write_text(content, encoding="utf-8")
     logger.debug("%s created.", orm_file)

datafaker/make.py

@@ -10,6 +10,7 @@
 from typing import Any, Final, Optional, Tuple, Type
 
 import pandas as pd
+import snowflake.sqlalchemy.custom_types as snowflake_types
 import snsql
 import typer
 import yaml
@@ -58,6 +59,8 @@
     if issubclass(entry, BaseProvider) and entry.__module__ == "datafaker.providers":
         PROVIDER_IMPORTS.append(entry_name)
 
+DEFAULT_UNIQUE_CONSTRAINT_TRIES: int | None = 500
+
 
 @dataclass
 class VocabularyTableGeneratorInfo:
@@ -378,6 +381,10 @@ def get_result_mappings(
         generator=_string_generator,
         choice=True,
     ),
+    snowflake_types.TIMESTAMP_LTZ: GeneratorInfo(
+        generator="generic.datetime.datetime",
+        arg_types={"start": int, "end": int},
+    ),
 }
 
 
@@ -654,6 +661,8 @@ def get_generation_info(
     max_unique_constraint_tries = get_property_or_none(
         config, "max-unique-constraint-tries", int
     )
+    if max_unique_constraint_tries is None:
+        max_unique_constraint_tries = DEFAULT_UNIQUE_CONSTRAINT_TRIES
     return GenerationInfo(
         provider_imports=PROVIDER_IMPORTS,
         row_generator_module_name=row_generator_module_name,
@@ -697,14 +706,17 @@ def _generate_vocabulary_table(
 def make_tables_file(
     db_dsn: str,
     schema_name: Optional[str],
+    views_to_tables: bool,
     parquet_dir: Optional[Path] = None,
     engine: Optional[Engine] = None,
 ) -> str:
     """Construct the YAML file representing the schema."""
     if engine is None:
         engine = get_sync_engine(create_db_engine(db_dsn, schema_name=schema_name))
 
-    metadata = get_metadata(engine, schema_name=schema_name)
+    metadata = get_metadata(
+        engine, schema_name=schema_name, views_to_tables=views_to_tables
+    )
     meta_dict = metadata_to_dict(metadata, schema_name, engine, parquet_dir)
 
     if parquet_dir is not None:

datafaker/proposers/__init__.py

@@ -17,7 +17,7 @@
     MultivariateNormalProposerFactory,
 )
 from datafaker.proposers.extract import DateComponentExtractProposerFactory
-from datafaker.proposers.intervals import DateAfterProposerFactory
+from datafaker.proposers.intervals import CopyProposerFactory, DateAfterProposerFactory
 from datafaker.proposers.mimesis import (
     MimesisDateProposerFactory,
     MimesisDateTimeProposerFactory,
@@ -57,6 +57,7 @@ def everything_factory(config: Mapping, metadata: MetaData) -> ProposerFactory:
         NullPartitionedNormalProposerFactory(config, metadata),
         NullPartitionedLogNormalProposerFactory(config, metadata),
         DateAfterProposerFactory(config, metadata),
+        CopyProposerFactory(config, metadata),
         DateComponentExtractProposerFactory(config, metadata),
         IncrementProposerFactory(),
     )