YAML vocab files

.gitignore

@@ -139,3 +139,9 @@ sqlsynthgen/config.ini
 # sphinx
 docs/build/*
 docs/temp/*
+
+# vim swap files
+*.swp
+
+ssg.py
+orm.py

sqlsynthgen/base.py

@@ -1,9 +1,9 @@
 """Base generator classes."""
-import csv
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
+import yaml
 from sqlalchemy import insert
 
 
@@ -15,9 +15,9 @@ class FileUploader:
 
     def load(self, connection: Any) -> None:
         """Load the data from file."""
-        with Path(self.table.fullname + ".csv").open(
+        with Path(self.table.fullname + ".yaml").open(
             "r", newline="", encoding="utf-8"
-        ) as csvfile:
-            reader = csv.DictReader(csvfile)
-            stmt = insert(self.table).values(list(reader))
+        ) as yamlfile:
+            rows = yaml.load(yamlfile, Loader=yaml.Loader)
+            stmt = insert(self.table).values(list(rows))
         connection.execute(stmt)

sqlsynthgen/create.py

@@ -1,7 +1,9 @@
 """Functions and classes to create and populate the target database."""
+import logging
 from typing import Any, Dict
 
 from sqlalchemy import create_engine, insert
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.schema import CreateSchema
 
 from sqlsynthgen.settings import get_settings
@@ -42,7 +44,11 @@ def create_db_vocab(vocab_dict: Dict[str, Any]) -> None:
 
     with dst_engine.connect() as dst_conn:
         for vocab_table in vocab_dict.values():
-            vocab_table.load(dst_conn)
+            try:
+                vocab_table.load(dst_conn)
+            except IntegrityError as e:
+                logging.warning("Loading the vocabulary table %s failed:", vocab_table)
+                logging.warning(e)
 
 
 def create_db_data(sorted_tables: list, generator_dict: dict, num_passes: int) -> None:

sqlsynthgen/main.py

@@ -114,6 +114,7 @@ def make_generators(
           Must be in the current working directory.
         ssg_file (str): Path to write the generators file to.
         config_file (str): Path to configuration file.
+        stats_file (str): Path to source stats file (output of make-stats).
     """
     ssg_file_path = Path(ssg_file)
     if ssg_file_path.exists():

sqlsynthgen/make.py

@@ -22,7 +22,7 @@
         "from mimesis.locales import Locale",
         "from sqlsynthgen.base import FileUploader",
         "",
-        "generic = Generic(locale=Locale.EN)",
+        "generic = Generic(locale=Locale.EN_GB)",
         "",
     )
 )

sqlsynthgen/utils.py

@@ -1,5 +1,4 @@
 """Utility functions."""
-import csv
 import os
 import sys
 from importlib import import_module
@@ -44,22 +43,19 @@ def import_file(file_name: str) -> ModuleType:
 
 
 def download_table(table: Any, engine: Any) -> None:
-    """Download a Table and store it as a .csv file."""
-    csv_file_name = table.fullname + ".csv"
-    csv_file_path = Path(csv_file_name)
-    if csv_file_path.exists():
-        print(f"{str(csv_file_name)} already exists. Exiting...", file=stderr)
+    """Download a Table and store it as a .yaml file."""
+    yaml_file_name = table.fullname + ".yaml"
+    yaml_file_path = Path(yaml_file_name)
+    if yaml_file_path.exists():
+        print(f"{str(yaml_file_name)} already exists. Exiting...", file=stderr)
         sys.exit(1)
 
     stmt = select([table])
     with engine.connect() as conn:
-        result = list(conn.execute(stmt))
+        result = [dict(row.items()) for row in conn.execute(stmt)]
 
-    with csv_file_path.open("w", newline="", encoding="utf-8") as csvfile:
-        writer = csv.writer(csvfile, delimiter=",")
-        writer.writerow([x.name for x in table.columns])
-        for row in result:
-            writer.writerow(row)
+    with yaml_file_path.open("w", newline="", encoding="utf-8") as yamlfile:
+        yamlfile.write(yaml.dump(result))
 
 
 def create_engine_with_search_path(postgres_dsn: PostgresDsn, schema_name: str) -> Any:

tests/examples/basetable.yaml

@@ -0,0 +1,3 @@
+- id: 1
+- id: 2
+- id: 3

tests/examples/example_config.yaml

@@ -52,3 +52,7 @@ tables:
           - visit_duration_seconds
   concept:
     vocabulary_table: true
+  concept_type:
+    vocabulary_table: true
+  mitigation_type:
+    vocabulary_table: true

tests/examples/expected.yaml

@@ -0,0 +1 @@
+- id: 1

tests/examples/expected_ssg.py

@@ -3,7 +3,7 @@
 from mimesis.locales import Locale
 from sqlsynthgen.base import FileUploader
 
-generic = Generic(locale=Locale.EN)
+generic = Generic(locale=Locale.EN_GB)
 
 from sqlsynthgen.providers import BytesProvider
 generic.add_provider(BytesProvider)