Add mariadb support

Author: Iain-S

poetry.lock

@@ -24,6 +24,7 @@ jsonschema = "^4.17.3"
 sqlacodegen = "3.0.0rc1"
 asyncpg = "^0.27.0"
 greenlet = "^2.0.2"
+pymysql = "^1.1.0"
 
 [tool.poetry.group.dev.dependencies]
 isort = "^5.10.1"

pyproject.toml

@@ -16,7 +16,7 @@ def create_db_tables(metadata: Any) -> Any:
     """Create tables described by the sqlalchemy metadata object."""
     settings = get_settings()
 
-    engine = create_db_engine(settings.dst_postgres_dsn)  # type: ignore
+    engine = create_db_engine(settings.dst_dsn)  # type: ignore
 
     # Create schema, if necessary.
     if settings.dst_schema:
@@ -26,7 +26,7 @@ def create_db_tables(metadata: Any) -> Any:
 
         # Recreate the engine, this time with a schema specified
         engine = create_db_engine(
-            settings.dst_postgres_dsn, schema_name=schema_name  # type: ignore
+            settings.dst_dsn, schema_name=schema_name  # type: ignore
         )
 
     metadata.create_all(engine)
@@ -37,7 +37,7 @@ def create_db_vocab(vocab_dict: Dict[str, Any]) -> None:
     settings = get_settings()
 
     dst_engine = create_db_engine(
-        settings.dst_postgres_dsn, schema_name=settings.dst_schema  # type: ignore
+        settings.dst_dsn, schema_name=settings.dst_schema  # type: ignore
     )
 
     with dst_engine.connect() as dst_conn:
@@ -60,7 +60,7 @@ def create_db_data(
     settings = get_settings()
 
     dst_engine = create_db_engine(
-        settings.dst_postgres_dsn, schema_name=settings.dst_schema  # type: ignore
+        settings.dst_dsn, schema_name=settings.dst_schema  # type: ignore
     )
 
     with dst_engine.connect() as dst_conn:

sqlsynthgen/create.py

@@ -10,7 +10,6 @@
 import yaml
 from jsonschema.exceptions import ValidationError
 from jsonschema.validators import validate
-from pydantic import PostgresDsn
 
 from sqlsynthgen.create import create_db_data, create_db_tables, create_db_vocab
 from sqlsynthgen.make import make_src_stats, make_table_generators, make_tables_file
@@ -35,12 +34,12 @@ def _check_file_non_existence(file_path: Path) -> None:
         sys.exit(1)
 
 
-def _get_src_postgres_dsn(settings: Settings) -> PostgresDsn:
+def _get_src_postgres_dsn(settings: Settings) -> str:
     """Return the source DB Postgres DSN.
 
     Check that source db details have been set. Exit with error message if not.
     """
-    if (src_dsn := settings.src_postgres_dsn) is None:
+    if (src_dsn := settings.src_dsn) is None:
         typer.echo("Missing source database connection details.", err=True)
         sys.exit(1)
     return src_dsn
@@ -180,7 +179,7 @@ def make_stats(
     config = read_yaml_file(config_file) if config_file is not None else {}
 
     settings = get_settings()
-    src_dsn: PostgresDsn = _get_src_postgres_dsn(settings)
+    src_dsn: str = _get_src_postgres_dsn(settings)
 
     src_stats = asyncio.get_event_loop().run_until_complete(
         make_src_stats(src_dsn, config, settings.src_schema)
@@ -211,7 +210,7 @@ def make_tables(
         _check_file_non_existence(orm_file_path)
 
     settings = get_settings()
-    src_dsn: PostgresDsn = _get_src_postgres_dsn(settings)
+    src_dsn: str = _get_src_postgres_dsn(settings)
 
     content = make_tables_file(src_dsn, settings.src_schema)
     orm_file_path.write_text(content, encoding="utf-8")

sqlsynthgen/main.py

@@ -13,7 +13,6 @@
 from black import FileMode, format_str
 from jinja2 import Environment, FileSystemLoader, Template
 from mimesis.providers.base import BaseProvider
-from pydantic import PostgresDsn
 from sqlacodegen.generators import DeclarativeGenerator
 from sqlalchemy import MetaData, UniqueConstraint, text
 from sqlalchemy.sql import sqltypes
@@ -351,7 +350,7 @@ def make_table_generators(
 
     settings = get_settings()
     engine = create_db_engine(
-        settings.src_postgres_dsn, schema_name=settings.src_schema  # type: ignore
+        settings.src_dsn, schema_name=settings.src_schema  # type: ignore
     )
 
     tables: List[TableGenerator] = []
@@ -430,7 +429,7 @@ def _get_generator_for_vocabulary_table(
     )
 
 
-def make_tables_file(db_dsn: PostgresDsn, schema_name: Optional[str]) -> str:
+def make_tables_file(db_dsn: str, schema_name: Optional[str]) -> str:
     """Write a file with the SQLAlchemy ORM classes.
 
     Exists with an error if sqlacodegen is unsuccessful.
@@ -455,7 +454,7 @@ def make_tables_file(db_dsn: PostgresDsn, schema_name: Optional[str]) -> str:
 
 
 async def make_src_stats(
-    dsn: PostgresDsn, config: dict, schema_name: Optional[str] = None
+    dsn: str, config: dict, schema_name: Optional[str] = None
 ) -> dict:
     """Run the src-stats queries specified by the configuration.
 

sqlsynthgen/make.py

@@ -11,10 +11,8 @@ def remove_db_data(orm_module: ModuleType, ssg_module: ModuleType) -> None:
     """Truncate the synthetic data tables but not the vocabularies."""
     settings = get_settings()
 
-    assert settings.dst_postgres_dsn, "Missing destination database settings"
-    dst_engine = create_db_engine(
-        settings.dst_postgres_dsn, schema_name=settings.dst_schema
-    )
+    assert settings.dst_dsn, "Missing destination database settings"
+    dst_engine = create_db_engine(settings.dst_dsn, schema_name=settings.dst_schema)
 
     with dst_engine.connect() as dst_conn:
         for table in reversed(orm_module.Base.metadata.sorted_tables):
@@ -27,10 +25,8 @@ def remove_db_vocab(orm_module: ModuleType, ssg_module: ModuleType) -> None:
     """Truncate the vocabulary tables."""
     settings = get_settings()
 
-    assert settings.dst_postgres_dsn, "Missing destination database settings"
-    dst_engine = create_db_engine(
-        settings.dst_postgres_dsn, schema_name=settings.dst_schema
-    )
+    assert settings.dst_dsn, "Missing destination database settings"
+    dst_engine = create_db_engine(settings.dst_dsn, schema_name=settings.dst_schema)
 
     with dst_engine.connect() as dst_conn:
         for table in reversed(orm_module.Base.metadata.sorted_tables):
@@ -43,10 +39,8 @@ def remove_db_tables(orm_module: ModuleType) -> None:
     """Drop the tables in the destination schema."""
     settings = get_settings()
 
-    assert settings.dst_postgres_dsn, "Missing destination database settings"
-    dst_engine = create_db_engine(
-        settings.dst_postgres_dsn, schema_name=settings.dst_schema
-    )
+    assert settings.dst_dsn, "Missing destination database settings"
+    dst_engine = create_db_engine(settings.dst_dsn, schema_name=settings.dst_schema)
 
     metadata = orm_module.Base.metadata
     metadata.drop_all(dst_engine)

sqlsynthgen/remove.py

@@ -19,7 +19,7 @@
 from typing import Any, Optional
 
 # pylint: disable=no-self-argument
-from pydantic import BaseSettings, PostgresDsn, validator
+from pydantic import BaseSettings, validator
 
 
 class Settings(BaseSettings):
@@ -31,96 +31,40 @@ class Settings(BaseSettings):
     and synthetic values inserted.
 
     Attributes:
-        src_host_name (str):
-            An element (host-name) of connection parameter
-        src_port (int):
-            Connection port eg. 5432
-        src_user_name (str) :
-            Connection username e.g. `postgres` or `myuser@mydb`
-        src_password (str) :
-            Connection password
-        src_db_name (str) :
-            Connection database e.g. "postgres"
-        src_ssl_required (bool) :
-            Flag `True` if db requires SSL
-
-        dst_host_name (str):
-            Connection host-name to destination db
-        dst_port (int) :
-            Connection port eg. 5432
-        dst_user_name (str) :
-            Connection username e.g. `postgres` or `myuser@mydb`
-        dst_password (str) :
-            Connection password
-        dst_db_name (str) :
-            Connection database e.g. `postgres`
-        dst_ssl_required (bool) :
-            Flag `True` if db requires SSL
+        src_dsn (str) :
+            A DSN for connecting to the source database.
+
+        src_schema (str) :
+            The source database schema to use, if applicable.
+
+        dst_dsn (str) :
+            A DSN for connecting to the destination database.
+
+        dst_schema (str) :
+            The destination database schema to use, if applicable.
     """
 
-    # Connection parameters for the source PostgreSQL database. See also
+    # Connection parameters for the source database. See also
     # https://www.postgresql.org/docs/11/libpq-connect.html#LIBPQ-PARAMKEYWORDS
-    src_host_name: Optional[str]  # e.g. "mydb.mydomain.com" or "0.0.0.0"
-    src_port: int = 5432
-    src_user_name: Optional[str]  # e.g. "postgres" or "myuser@mydb"
-    src_password: Optional[str]
-    src_db_name: Optional[str]
-    src_ssl_required: bool = False  # whether the db requires SSL
-    src_schema: Optional[str]
+    src_dsn: Optional[str]
+    dst_dsn: Optional[str]
 
-    # Connection parameters for the destination PostgreSQL database.
-    dst_host_name: Optional[
-        str
-    ]  # Connection parameter e.g. "mydb.mydomain.com" or "0.0.0.0"
-    dst_port: int = 5432
-    dst_user_name: Optional[str]  # e.g. "postgres" or "myuser@mydb"
-    dst_password: Optional[str]
-    dst_db_name: Optional[str]
+    src_schema: Optional[str]
     dst_schema: Optional[str]
-    dst_ssl_required: bool = False  # whether the db requires SSL
-
-    # These are calculated so do not provide them explicitly
-    src_postgres_dsn: Optional[PostgresDsn]
-    dst_postgres_dsn: Optional[PostgresDsn]
-
-    @validator("src_postgres_dsn", pre=True)
-    def validate_src_postgres_dsn(
-        cls, _: Optional[PostgresDsn], values: Any
-    ) -> Optional[str]:
-        """Create and validate the source db data source name."""
-        return cls.check_postgres_dsn(_, values, "src")
-
-    @validator("dst_postgres_dsn", pre=True)
-    def validate_dst_postgres_dsn(
-        cls, _: Optional[PostgresDsn], values: Any
-    ) -> Optional[str]:
-        """Create and validate the destination db data source name."""
-        return cls.check_postgres_dsn(_, values, "dst")
-
-    @staticmethod
-    def check_postgres_dsn(
-        _: Optional[PostgresDsn], values: Any, prefix: str
-    ) -> Optional[str]:
-        """Build a DSN string from the host, db name, port, username and password."""
-        # We want to build the Data Source Name ourselves so none should be provided
-        if _:
-            raise ValueError("postgres_dsn should not be provided")
-
-        user = values[f"{prefix}_user_name"]
-        password = values[f"{prefix}_password"]
-        host = values[f"{prefix}_host_name"]
-        port = values[f"{prefix}_port"]
-        db_name = values[f"{prefix}_db_name"]
-
-        if user and password and host and port and db_name:
-            dsn = f"postgresql://{user}:{password}@{host}:{port}/{db_name}"
-
-            if values[f"{prefix}_ssl_required"]:
-                return dsn + "?sslmode=require"
-
-            return dsn
-
-        return None
+
+    @validator("src_dsn")
+    def validate_src_dsn(cls, dsn: Optional[str], values: Any) -> Optional[str]:
+        """Create and validate the source DB DSN."""
+        if dsn and dsn.startswith("mariadb"):
+            assert values.get("src_schema") is None
+        return dsn
+
+    @validator("dst_dsn")
+    def validate_dst_dsn(cls, dsn: Optional[str], values: Any) -> Optional[str]:
+        """Create and validate the destination DB DSN."""
+        if dsn and dsn.startswith("mariadb"):
+            assert values.get("dst_schema") is None
+        return dsn
 
     @dataclass
     class Config:

sqlsynthgen/settings.py

@@ -7,7 +7,6 @@
 from typing import Any, Optional
 
 import yaml
-from pydantic import PostgresDsn
 from sqlalchemy import create_engine, event, select
 from sqlalchemy.ext.asyncio import create_async_engine
 
@@ -53,7 +52,7 @@ def download_table(table: Any, engine: Any, yaml_file_name: str) -> None:
 
 
 def create_db_engine(
-    postgres_dsn: PostgresDsn,
+    postgres_dsn: str,
     schema_name: Optional[str] = None,
     use_asyncio: bool = False,
     **kwargs: dict,

sqlsynthgen/utils.py

@@ -44,7 +44,7 @@ def test_create_db_tables(
 
         create_db_tables(mock_meta)
         mock_create_engine.assert_called_once_with(
-            mock_get_settings.return_value.dst_postgres_dsn
+            mock_get_settings.return_value.dst_dsn
         )
         mock_meta.create_all.assert_called_once_with(mock_create_engine.return_value)
 
@@ -127,7 +127,7 @@ def test_create_db_vocab(
             mock_create_engine.return_value.connect.return_value.__enter__.return_value
         )
         mock_create_engine.assert_called_once_with(
-            mock_get_settings.return_value.dst_postgres_dsn
+            mock_get_settings.return_value.dst_dsn
         )
         # Running the same insert twice should be fine.
         create_db_vocab(vocab_list)

tests/test_create.py

@@ -274,7 +274,7 @@ def test_make_tables_with_force_enabled(
                 result: Result = runner.invoke(app, ["make-tables", force_option])
 
                 mock_make_tables.assert_called_once_with(
-                    test_settings.src_postgres_dsn, test_settings.src_schema
+                    test_settings.src_dsn, test_settings.src_schema
                 )
                 mock_path.return_value.write_text.assert_called_once_with(
                     mock_tables_output, encoding="utf-8"
@@ -308,9 +308,7 @@ def test_make_stats(
         self.assertSuccess(result)
         with open(example_conf_path, "r", encoding="utf8") as f:
             config = yaml.safe_load(f)
-        mock_make.assert_called_once_with(
-            get_test_settings().src_postgres_dsn, config, None
-        )
+        mock_make.assert_called_once_with(get_test_settings().src_dsn, config, None)
         mock_path.return_value.write_text.assert_called_once_with(
             "a: 1\n", encoding="utf-8"
         )
@@ -391,7 +389,7 @@ def test_make_stats_with_force_enabled(
                 )
 
                 mock_make.assert_called_once_with(
-                    test_settings.src_postgres_dsn, config_file_content, None
+                    test_settings.src_dsn, config_file_content, None
                 )
                 mock_path.return_value.write_text.assert_called_once_with(
                     "some_stat: 0\n", encoding="utf-8"