Add option to not use SmartnoiseSQL

pyproject.toml

@@ -18,7 +18,7 @@ pyyaml = "^5.0"
 sqlalchemy = "^1.4"
 sphinx-rtd-theme = {version = "^1.2.0", optional = true}
 sphinxcontrib-napoleon = {version = "^0.7", optional = true}
-smartnoise-sql = "^0.2.9.1"
+smartnoise-sql = "^0.2.11"
 
 [tool.poetry.group.dev.dependencies]
 black = "^22.10.0"

sqlsynthgen/make.py

@@ -2,13 +2,13 @@
 import inspect
 from sys import stderr
 from types import ModuleType
-from typing import Any, Final, Optional
+from typing import Any, Dict, Final, Optional
 
 import snsql
 from mimesis.providers.base import BaseProvider
 from pydantic import PostgresDsn
 from sqlacodegen.generators import DeclarativeGenerator
-from sqlalchemy import MetaData, create_engine
+from sqlalchemy import MetaData, create_engine, text
 from sqlalchemy.sql import sqltypes
 
 from sqlsynthgen import providers
@@ -256,12 +256,15 @@ def make_src_stats(
     else:
         engine = create_engine(dsn, echo=False, future=True)
 
-    dp_config = config.get("smartnoise-sql", {})
-    snsql_metadata = {"": dp_config}
-    src_stats = {}
-    for stat_data in config.get("src-stats", []):
-        privacy = snsql.Privacy(epsilon=stat_data["epsilon"], delta=stat_data["delta"])
-        with engine.connect() as conn:
+    use_smartnoise_sql = config.get("use-smartnoise-sql", True)
+    if use_smartnoise_sql:
+        dp_config = config.get("smartnoise-sql", {})
+        snsql_metadata = {"": dp_config}
+
+        def execute_query(conn: Any, stat_data: Dict[str, Any]) -> Any:
+            privacy = snsql.Privacy(
+                epsilon=stat_data["epsilon"], delta=stat_data["delta"]
+            )
             reader = snsql.from_connection(
                 conn.connection,
                 engine="postgres",
@@ -270,5 +273,17 @@ def make_src_stats(
             )
             private_result = reader.execute(stat_data["query"])
             # The first entry in the list names the columns, skip that.
-            src_stats[stat_data["name"]] = private_result[1:]
+            return private_result[1:]
+
+    else:
+
+        def execute_query(conn: Any, stat_data: Dict[str, Any]) -> Any:
+            result = conn.execute(text(stat_data["query"])).fetchall()
+            return [list(r) for r in result]
+
+    with engine.connect() as conn:
+        src_stats = {
+            stat_data["name"]: execute_query(conn, stat_data)
+            for stat_data in config.get("src-stats", [])
+        }
     return src_stats

tests/test_make.py

@@ -138,11 +138,13 @@ def test_make_stats(self) -> None:
         conf_path = Path("example_config.yaml")
         with open(conf_path, "r", encoding="utf8") as f:
             config = yaml.safe_load(f)
+        config_no_snsql = {**config, "use-smartnoise-sql": False}
 
         # Check that make_src_stats works with, or without, a schema
         for args in (
             (connection_string, config),
             (connection_string, config, "public"),
+            (connection_string, config_no_snsql),
         ):
 
             src_stats = make_src_stats(*args)