Initial Snowflake support

.github/workflows/docs.yml

@@ -4,7 +4,7 @@ on:
   push:
     branches:
       # need this to be able to run the workflow until it has been merged into main
-      - yhong123/131-update_documentation
+      - yhong123/112-column_evaluators
   workflow_dispatch:
 env:
   PYTHON_VERSION: "3.12"

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,
@@ -168,22 +160,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:
@@ -226,11 +230,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
@@ -266,16 +274,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
@@ -298,25 +304,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 fk_refers_to_ignored_table(fk: ForeignKey) -> bool:
     """
     Test if this foreign key refers to an ignored table.

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
@@ -211,6 +212,96 @@ def compile_stddev_mssql(element: StdDev, compiler: Any, **kw: Any) -> str:
     return f"STDEV({e})"
 
 
+class WordCount(ColumnElement[int]):  # pylint: disable=too-many-ancestors
+    """Count whitespace-separated words in a text expression."""
+
+    expr: ColumnElement[str]
+
+    _traverse_internals = [
+        ("expr", InternalTraversal.dp_clauseelement),
+    ]
+
+    def __init__(self, expr: ColumnElement[str]):
+        """Get a clause for the number of whitespace-separated words in ``expr``."""
+        self.expr = expr
+
+    __sa_operate__ = ColumnElement.operate
+
+
+@compiles(WordCount)
+def compile_word_count(element: WordCount, compiler: Any, **kw: Any) -> str:
+    """Exact word count via regexp_split_to_array (Postgres/DuckDB)."""
+    e = compiler.process(element.expr, **kw)
+    return f"array_length(regexp_split_to_array(trim({e}), '\\s+'), 1)"
+
+
+@compiles(WordCount, "mssql")
+def compile_word_count_mssql(element: WordCount, compiler: Any, **kw: Any) -> str:
+    """MSSQL equivalent: approximate word count via space-counting.
+
+    No regex/array-split support, so approximate word count as (space
+    count in the trimmed text) + 1. This overcounts by one per run of 2+
+    consecutive spaces relative to the exact regex-based count used on
+    Postgres/DuckDB, but there is no T-SQL primitive for a real
+    whitespace-run split.
+    """
+    e = compiler.process(element.expr, **kw)
+    trimmed = f"TRIM({e})"
+    return (
+        f"(CASE WHEN LEN({trimmed}) = 0 THEN 0 ELSE "
+        f"LEN({trimmed}) - LEN(REPLACE({trimmed}, ' ', '')) + 1 END)"
+    )
+
+
+class SentenceCount(ColumnElement[int]):  # pylint: disable=too-many-ancestors
+    """Count '.'/'!'/'?'-delimited sentences in a text expression."""
+
+    expr: ColumnElement[str]
+
+    _traverse_internals = [
+        ("expr", InternalTraversal.dp_clauseelement),
+    ]
+
+    def __init__(self, expr: ColumnElement[str]):
+        """Get a clause for the number of sentences in ``expr``."""
+        self.expr = expr
+
+    __sa_operate__ = ColumnElement.operate
+
+
+@compiles(SentenceCount)
+def compile_sentence_count(element: SentenceCount, compiler: Any, **kw: Any) -> str:
+    """Exact sentence count via regexp_split_to_array (Postgres/DuckDB)."""
+    e = compiler.process(element.expr, **kw)
+    return (
+        f"array_length(regexp_split_to_array(trim({e}, ' .!?\\t\\n\\r'), '[.!?]+'), 1)"
+    )
+
+
+@compiles(SentenceCount, "mssql")
+def compile_sentence_count_mssql(
+    element: SentenceCount, compiler: Any, **kw: Any
+) -> str:
+    """MSSQL equivalent: approximate sentence count via delimiter-counting.
+
+    No regex support, so approximate sentence count by normalizing '!'/'?'
+    to '.' and counting '.' characters, after trimming whitespace and
+    sentence-enders from both ends (so a leading/trailing sentence-ender
+    doesn't add a spurious segment, matching the Postgres/DuckDB
+    behaviour). This overcounts by one per run of 2+ consecutive delimiter
+    characters (e.g. "Wow!!") relative to the exact regex-based count, but
+    there is no T-SQL primitive for regex-run splitting.
+    """
+    e = compiler.process(element.expr, **kw)
+    trim_chars = "' .!?' + CHAR(9) + CHAR(10) + CHAR(13)"
+    trimmed = f"TRIM({trim_chars} FROM {e})"
+    normalized = f"REPLACE(REPLACE({trimmed}, '!', '.'), '?', '.')"
+    return (
+        f"(CASE WHEN LEN({trimmed}) = 0 THEN 0 ELSE "
+        f"LEN({normalized}) - LEN(REPLACE({normalized}, '.', '')) + 1 END)"
+    )
+
+
 class IsNull(ColumnElement[bool]):  # pylint: disable=too-many-ancestors
     """Represent IS NULL as an expression."""
 
@@ -425,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)