Fix remaining RANDOM()/LIMIT dialect incompatibilities (closes #107)

Six locations in generators and the interactive shell still emitted
RANDOM() and LIMIT which MS-SQL does not support.

Group A — interactive shell (base.py, table.py, generators.py):
  Replace sqlalchemy.text() f-strings with SQLAlchemy expression API.
  func.newid() vs func.random() chosen per dialect.name == "mssql";
  .limit(n) compiles to TOP n / LIMIT n automatically.

Group B — CovariateQuery (continuous.py):
  Add dialect_name parameter to __init__(); _inner_query() now emits
  SELECT TOP n … ORDER BY NEWID() on MS-SQL and SELECT * … ORDER BY
  RANDOM() LIMIT n elsewhere. dialect_name passed from get_generators().

Group C — MissingnessType (missingness.py):
  sampled_query() accepts dialect_name; returns the MS-SQL TOP/NEWID
  form when dialect_name == "mssql". do_sampled() passes
  self.sync_engine.dialect.name.

Tests added in test_generators_dialect.py (CovariateQuery, Missingness)
and new test_interactive_dialect.py (peek, _get_column_data,
print_column_data) covering both MS-SQL and PostgreSQL.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: myyong

datafaker/generators/continuous.py

@@ -396,13 +396,15 @@ def __init__(
         self,
         table: str,
         factory: MultivariateNormalGeneratorFactoryBase,
+        dialect_name: str = "",
     ) -> None:
         """
         Initialize the query for the basics for multivariate normal/lognormal parameters.
 
         :param table: The name of the table to be queried.
         :param factory: The generator factory, perhaps with overridden
         ``query_var`` and ``query_predicate`` methods.
+        :param dialect_name: The SQLAlchemy dialect name (e.g. ``"mssql"``).
         """
         self.table = table
         self._columns: Sequence[Column] = []
@@ -412,6 +414,7 @@ def __init__(
         self.suppress_count = 1
         self._sample_count: int | None = None
         self._factory = factory
+        self._dialect_name = dialect_name
         self._predicate_fn = lambda x: x + " IS NOT NULL"
 
     def get_query_comment(self) -> str:
@@ -566,6 +569,11 @@ def _inner_query(self) -> str:
             where = " WHERE " + where
         if self._sample_count is None:
             return self.table + where
+        if self._dialect_name == "mssql":
+            return (
+                f"(SELECT TOP {self._sample_count} * FROM {self.table}{where}"
+                f" ORDER BY NEWID()) AS _sampled"
+            )
         return (
             f"(SELECT * FROM {self.table}{where} ORDER BY RANDOM()"
             f" LIMIT {self._sample_count}) AS _sampled"
@@ -628,7 +636,7 @@ def get_generators(
                 return []
         column_names = [c.name for c in columns]
         table = columns[0].table.name
-        cq = CovariateQuery(table, self).columns(columns)
+        cq = CovariateQuery(table, self, dialect_name=engine.dialect.name).columns(columns)
         query = cq.get()
         with engine.connect() as connection:
             try:

datafaker/interactive/base.py

@@ -10,7 +10,7 @@
 
 import sqlalchemy
 from prettytable import PrettyTable
-from sqlalchemy import Engine, ForeignKey, MetaData, Table
+from sqlalchemy import Engine, ForeignKey, MetaData, Table, func, literal_column, or_, select, table as sa_table
 from typing_extensions import Self
 
 from datafaker.utils import (
@@ -412,19 +412,23 @@ def do_peek(self, arg: str) -> None:
         col_names = arg.split()
         if not col_names:
             col_names = self._get_column_names()
-        nonnulls = [f'"{cn}" IS NOT NULL' for cn in col_names]
+        random_fn = (
+            func.newid() if self.sync_engine.dialect.name == "mssql" else func.random()
+        )
+        col_exprs = [literal_column(f'"{cn}"') for cn in col_names]
+        nonnull_clauses = [literal_column(f'"{cn}"').isnot(None) for cn in col_names]
+        stmt = (
+            select(*col_exprs)
+            .select_from(sa_table(table_name))
+            .where(or_(*nonnull_clauses))
+            .order_by(random_fn)
+            .limit(max_peek_rows)
+        )
         with self.sync_engine.connect() as connection:
-            cols = ", ".join(f'"{cn}"' for cn in col_names)
-            where = "WHERE" if nonnulls else ""
-            nonnull = " OR ".join(nonnulls)
-            query = sqlalchemy.text(
-                f"SELECT {cols} FROM {table_name} {where} {nonnull}"
-                f" ORDER BY RANDOM() LIMIT {max_peek_rows}"
-            )
             try:
-                result = connection.execute(query)
+                result = connection.execute(stmt)
             except sqlalchemy.exc.SQLAlchemyError as exc:
-                self.print(self.ERROR_FAILED_SQL, exc=exc, query=query)
+                self.print(self.ERROR_FAILED_SQL, exc=exc, query=stmt)
                 return
             self.print_table(list(result.keys()), result.fetchmany(max_peek_rows))
 

datafaker/interactive/generators.py

@@ -6,7 +6,7 @@
 from typing import Any, Callable, Optional, cast
 
 import sqlalchemy
-from sqlalchemy import Column
+from sqlalchemy import Column, and_, func, literal_column, select, table as sa_table
 
 from datafaker.generators import everything_factory
 from datafaker.generators.base import Generator, PredefinedGenerator
@@ -754,15 +754,20 @@ def _get_column_data(
         self, count: int, to_str: Callable[[Any], str] = repr
     ) -> list[list[str]]:
         columns = self._get_column_names()
-        columns_string = ", ".join(columns)
-        pred = " AND ".join(f"{column} IS NOT NULL" for column in columns)
+        random_fn = (
+            func.newid() if self.sync_engine.dialect.name == "mssql" else func.random()
+        )
+        col_exprs = [literal_column(col) for col in columns]
+        nonnull_clauses = [literal_column(col).isnot(None) for col in columns]
+        stmt = (
+            select(*col_exprs)
+            .select_from(sa_table(self.table_name()))
+            .where(and_(*nonnull_clauses))
+            .order_by(random_fn)
+            .limit(count)
+        )
         with self.sync_engine.connect() as connection:
-            result = connection.execute(
-                sqlalchemy.text(
-                    f"SELECT {columns_string} FROM {self.table_name()}"
-                    f" WHERE {pred} ORDER BY RANDOM() LIMIT {count}"
-                )
-            )
+            result = connection.execute(stmt)
             return [[to_str(x) for x in xs] for xs in result.all()]
 
     def do_propose(self, _arg: str) -> None:

datafaker/interactive/missingness.py

@@ -23,19 +23,32 @@ class MissingnessType:
     columns: list[str]
 
     @classmethod
-    def sampled_query(cls, table: str, count: int, column_names: Iterable[str]) -> str:
+    def sampled_query(
+        cls,
+        table: str,
+        count: int,
+        column_names: Iterable[str],
+        dialect_name: str = "",
+    ) -> str:
         """
         Construct a query to make a sampling of the named rows of the table.
 
         :param table: The name of the table to sample.
         :param count: The number of samples to get.
         :param column_names: The columns to fetch.
+        :param dialect_name: The SQLAlchemy dialect name (e.g. ``"mssql"``).
         :return: The SQL query to do the sampling.
         """
         result_names = ", ".join([f"{c}__is_null" for c in column_names])
         column_is_nulls = ", ".join(
             [f"{c} IS NULL AS {c}__is_null" for c in column_names]
         )
+        if dialect_name == "mssql":
+            return (
+                f"SELECT COUNT(*) AS row_count, {result_names} FROM "
+                f"(SELECT TOP {count} {column_is_nulls} FROM {table} ORDER BY NEWID())"
+                f" AS __t GROUP BY {result_names}"
+            )
         return cls.SAMPLED_QUERY.format(
             result_names=result_names,
             column_is_nulls=column_is_nulls,
@@ -330,6 +343,7 @@ def do_sampled(self, arg: str) -> None:
                 entry.name,
                 count,
                 self.get_nullable_columns(entry.name),
+                dialect_name=self.sync_engine.dialect.name,
             ),
             [
                 "The missingness patterns and how often they appear in a"

datafaker/interactive/table.py

@@ -4,6 +4,7 @@
 from typing import Any, cast
 
 import sqlalchemy
+from sqlalchemy import func, literal_column, select, table as sa_table, text
 
 from datafaker.interactive.base import (
     TYPE_LETTER,
@@ -477,16 +478,23 @@ def print_column_data(self, column: str, count: int, min_length: int) -> None:
         :param count: The number of rows to sample.
         :param min_length: The minimum length of text to choose from (0 for any text).
         """
-        where = f"WHERE {column} IS NOT NULL"
+        random_fn = (
+            func.newid() if self.sync_engine.dialect.name == "mssql" else func.random()
+        )
+        col_expr = literal_column(column)
         if 0 < min_length:
-            where = f"WHERE LENGTH({column}) >= {min_length}"
+            where_clause = func.length(col_expr) >= min_length
+        else:
+            where_clause = col_expr.isnot(None)
+        stmt = (
+            select(col_expr)
+            .select_from(sa_table(self.table_name()))
+            .where(where_clause)
+            .order_by(random_fn)
+            .limit(count)
+        )
         with self.sync_engine.connect() as connection:
-            result = connection.execute(
-                sqlalchemy.text(
-                    f"SELECT {column} FROM {self.table_name()}"
-                    f" {where} ORDER BY RANDOM() LIMIT {count}"
-                )
-            )
+            result = connection.execute(stmt)
             self.columnize([str(x[0]) for x in result.all()])
 
     def print_row_data(self, count: int) -> None:
@@ -495,12 +503,17 @@ def print_row_data(self, count: int) -> None:
 
         :param count: The number of rows to report.
         """
+        random_fn = (
+            func.newid() if self.sync_engine.dialect.name == "mssql" else func.random()
+        )
+        stmt = (
+            select(text("*"))
+            .select_from(sa_table(self.table_name()))
+            .order_by(random_fn)
+            .limit(count)
+        )
         with self.sync_engine.connect() as connection:
-            result = connection.execute(
-                sqlalchemy.text(
-                    f"SELECT * FROM {self.table_name()} ORDER BY RANDOM() LIMIT {count}"
-                )
-            )
+            result = connection.execute(stmt)
             if result is None:
                 self.print("No rows in this table!")
                 return

tests/test_generators_dialect.py

@@ -309,3 +309,86 @@ def test_no_schema_omits_qualifier(self) -> None:
         """Without schema, the FROM clause has no dot-qualifier."""
         sql = self._get_make_buckets_sql("postgresql", schema=None)
         self.assertNotIn(".", sql)
+
+
+class TestCovariateQueryDialect(unittest.TestCase):
+    """CovariateQuery._inner_query() uses TOP/NEWID on MS-SQL and RANDOM/LIMIT elsewhere."""
+
+    def _make_factory(self) -> MagicMock:
+        factory = MagicMock()
+        factory.query_predicate.return_value = ""
+        return factory
+
+    def _inner_query(self, dialect_name: str) -> str:
+        from datafaker.generators.continuous import CovariateQuery
+
+        cq = (
+            CovariateQuery("person", self._make_factory(), dialect_name=dialect_name)
+            .sample_count(500)
+        )
+        return cq._inner_query().upper()
+
+    def test_mssql_uses_top_and_newid(self) -> None:
+        """MS-SQL inner query uses SELECT TOP n … ORDER BY NEWID()."""
+        sql = self._inner_query("mssql")
+        self.assertIn("TOP 500", sql)
+        self.assertIn("NEWID()", sql)
+        self.assertNotIn("RANDOM()", sql)
+        self.assertNotIn("LIMIT", sql)
+
+    def test_postgresql_uses_random_and_limit(self) -> None:
+        """PostgreSQL inner query uses ORDER BY RANDOM() LIMIT n."""
+        sql = self._inner_query("postgresql")
+        self.assertIn("RANDOM()", sql)
+        self.assertIn("LIMIT 500", sql)
+        self.assertNotIn("NEWID()", sql)
+        self.assertNotIn("TOP", sql)
+
+    def test_no_sample_count_has_no_random_or_limit(self) -> None:
+        """When sample_count is None no random ordering is emitted."""
+        from datafaker.generators.continuous import CovariateQuery
+
+        for dialect in ("mssql", "postgresql", ""):
+            with self.subTest(dialect=dialect):
+                cq = CovariateQuery("person", self._make_factory(), dialect_name=dialect)
+                sql = cq._inner_query().upper()
+                self.assertNotIn("RANDOM()", sql)
+                self.assertNotIn("NEWID()", sql)
+                self.assertNotIn("LIMIT", sql)
+                self.assertNotIn("TOP", sql)
+
+
+class TestMissingnessQueryDialect(unittest.TestCase):
+    """MissingnessType.sampled_query() produces dialect-correct SQL."""
+
+    def test_mssql_uses_top_and_newid(self) -> None:
+        """MS-SQL sampled query uses SELECT TOP n … ORDER BY NEWID()."""
+        from datafaker.interactive.missingness import MissingnessType
+
+        sql = MissingnessType.sampled_query(
+            "person", 1000, ["col_a", "col_b"], dialect_name="mssql"
+        ).upper()
+        self.assertIn("TOP 1000", sql)
+        self.assertIn("NEWID()", sql)
+        self.assertNotIn("RANDOM()", sql)
+        self.assertNotIn("LIMIT", sql)
+
+    def test_default_uses_random_and_limit(self) -> None:
+        """Default (no dialect) sampled query uses RANDOM() and LIMIT."""
+        from datafaker.interactive.missingness import MissingnessType
+
+        sql = MissingnessType.sampled_query("person", 1000, ["col_a"]).upper()
+        self.assertIn("RANDOM()", sql)
+        self.assertIn("LIMIT 1000", sql)
+        self.assertNotIn("NEWID()", sql)
+        self.assertNotIn("TOP", sql)
+
+    def test_mssql_result_contains_column_null_checks(self) -> None:
+        """MS-SQL sampled query retains IS NULL expressions for the named columns."""
+        from datafaker.interactive.missingness import MissingnessType
+
+        sql = MissingnessType.sampled_query(
+            "person", 500, ["gender_concept_id"], dialect_name="mssql"
+        )
+        self.assertIn("gender_concept_id IS NULL", sql)
+        self.assertIn("gender_concept_id__is_null", sql)