Add support for MySQL 8

DEVELOPING.md

@@ -53,7 +53,7 @@ one, run:
 
 ```bash
 export PGPASSWORD=dbos
-python3 dbos/_templates/hello/start_postgres_docker.py
+python3 dbos/_templates/dbos-db-starter/start_postgres_docker.py
 ```
 
 A successful test run results in the following output:

alembic.ini

@@ -3,6 +3,6 @@
 [alembic]
 # path to migration scripts
 # Use forward slashes (/) also on windows to provide an os agnostic path
-script_location = dbos/_migrations
+script_location = dbos/_migrations/mysql
 
 version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.

dbos/_app_db.py

@@ -1,13 +1,19 @@
 from typing import Optional, TypedDict
 
 import sqlalchemy as sa
+import sqlalchemy.dialects.mysql as mysql
 import sqlalchemy.dialects.postgresql as pg
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
 from ._dbos_config import ConfigFile
 from ._error import DBOSWorkflowConflictIDError
-from ._schemas.application_database import ApplicationSchema
+from ._logger import dbos_logger
+from ._schemas._mysql import Expressions
+from ._schemas.application_database import (
+    ApplicationSchema,
+    configure_application_schema_mysql,
+)
 
 
 class TransactionResultInternal(TypedDict):
@@ -29,37 +35,70 @@ class ApplicationDatabase:
 
     def __init__(self, config: ConfigFile):
         self.config = config
+        self.db_type = config["database"]["type"]
 
         app_db_name = config["database"]["app_db_name"]
+        if "postgresql" == self.db_type:
+            # If the application database does not already exist, create it
+            postgres_db_url = sa.URL.create(
+                "postgresql+psycopg",
+                username=config["database"]["username"],
+                password=config["database"]["password"],
+                host=config["database"]["hostname"],
+                port=config["database"]["port"],
+                database="postgres",
+            )
+            postgres_db_engine = sa.create_engine(postgres_db_url)
+            with postgres_db_engine.connect() as conn:
+                conn.execution_options(isolation_level="AUTOCOMMIT")
+                if not conn.execute(
+                    sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
+                    parameters={"db_name": app_db_name},
+                ).scalar():
+                    conn.execute(sa.text(f"CREATE DATABASE {app_db_name}"))
+            postgres_db_engine.dispose()
+
+            # Create a connection pool for the application database
+            app_db_url = sa.URL.create(
+                "postgresql+psycopg",
+                username=config["database"]["username"],
+                password=config["database"]["password"],
+                host=config["database"]["hostname"],
+                port=config["database"]["port"],
+                database=app_db_name,
+            )
+        elif "mysql" == self.db_type:
+            configure_application_schema_mysql(db_schema_name=app_db_name)
+            db_url_args = {
+                "drivername": "mysql+pymysql",
+                "username": config["database"]["username"],
+                "password": config["database"]["password"],
+                "host": config["database"]["hostname"],
+                "port": config["database"]["port"],
+            }
+            mysql_db_url = sa.URL.create(**db_url_args)
+            engine = sa.create_engine(mysql_db_url)
+            with engine.connect() as conn:
+                conn.execution_options(isolation_level="AUTOCOMMIT")
+                conn.execute(
+                    sa.text(
+                        f"""
+                                CREATE DATABASE IF NOT EXISTS `{app_db_name}`
+                                CHARACTER SET utf8mb4
+                                COLLATE utf8mb4_bin ;
+                            """
+                    )
+                )
+                dbos_logger.info(f"application database exists: {app_db_name}")
+            engine.dispose()
+
+            db_url_args["database"] = app_db_name
+            app_db_url = sa.URL.create(**db_url_args)
+        else:
+            raise RuntimeError(
+                f"unsupported database type: {config['database']['type']}"
+            )
 
-        # If the application database does not already exist, create it
-        postgres_db_url = sa.URL.create(
-            "postgresql+psycopg",
-            username=config["database"]["username"],
-            password=config["database"]["password"],
-            host=config["database"]["hostname"],
-            port=config["database"]["port"],
-            database="postgres",
-        )
-        postgres_db_engine = sa.create_engine(postgres_db_url)
-        with postgres_db_engine.connect() as conn:
-            conn.execution_options(isolation_level="AUTOCOMMIT")
-            if not conn.execute(
-                sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
-                parameters={"db_name": app_db_name},
-            ).scalar():
-                conn.execute(sa.text(f"CREATE DATABASE {app_db_name}"))
-        postgres_db_engine.dispose()
-
-        # Create a connection pool for the application database
-        app_db_url = sa.URL.create(
-            "postgresql+psycopg",
-            username=config["database"]["username"],
-            password=config["database"]["password"],
-            host=config["database"]["hostname"],
-            port=config["database"]["port"],
-            database=app_db_name,
-        )
         self.engine = sa.create_engine(
             app_db_url, pool_size=20, max_overflow=5, pool_timeout=30
         )
@@ -76,10 +115,22 @@ def __init__(self, config: ConfigFile):
     def destroy(self) -> None:
         self.engine.dispose()
 
-    @staticmethod
+    def _raise_unsupported_db_type(self):
+        raise RuntimeError(
+            f"unsupported database type: {self.db_type} (configured: {self.config['database']['type']})"
+        )
+
     def record_transaction_output(
-        session: Session, output: TransactionResultInternal
+        self, session: Session, output: TransactionResultInternal
     ) -> None:
+        if "postgresql" == self.db_type:
+            self._record_transaction_output_pg(session, output)
+        elif "mysql" == self.db_type:
+            self._record_transaction_output_mysql(session, output)
+        else:
+            self._raise_unsupported_db_type()
+
+    def _record_transaction_output_pg(self, session, output):
         try:
             session.execute(
                 pg.insert(ApplicationSchema.transaction_outputs).values(
@@ -99,7 +150,38 @@ def record_transaction_output(
                 raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
             raise
 
+    def _record_transaction_output_mysql(self, session, output):
+        try:
+            session.execute(
+                mysql.insert(ApplicationSchema.transaction_outputs).values(
+                    workflow_uuid=output["workflow_uuid"],
+                    function_id=output["function_id"],
+                    output=output["output"],
+                    error=None,
+                    txn_id=sa.text(
+                        "(SELECT TRX_ID FROM INFORMATION_SCHEMA.INNODB_TRX WHERE TRX_MYSQL_THREAD_ID = CONNECTION_ID())"
+                    ),
+                    txn_snapshot=output["txn_snapshot"],
+                    executor_id=(
+                        output["executor_id"] if output["executor_id"] else None
+                    ),
+                )
+            )
+        except DBAPIError as dbapi_error:
+            dbos_logger.warning(
+                f'error recording transaction output: {output["output"]}; dbapi_error: {dbapi_error}'
+            )
+            raise dbapi_error
+
     def record_transaction_error(self, output: TransactionResultInternal) -> None:
+        if "postgresql" == self.db_type:
+            self._record_transaction_error_pg(output)
+        elif "mysql" == self.db_type:
+            self._record_transaction_error_mysql(output)
+        else:
+            self._raise_unsupported_db_type()
+
+    def _record_transaction_error_pg(self, output):
         try:
             with self.engine.begin() as conn:
                 conn.execute(
@@ -122,6 +204,28 @@ def record_transaction_error(self, output: TransactionResultInternal) -> None:
                 raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
             raise
 
+    def _record_transaction_error_mysql(self, output):
+        try:
+            with self.engine.begin() as conn:
+                conn.execute(
+                    mysql.insert(ApplicationSchema.transaction_outputs).values(
+                        workflow_uuid=output["workflow_uuid"],
+                        function_id=output["function_id"],
+                        output=None,
+                        error=output["error"],
+                        txn_id=sa.text(Expressions.get_current_txid_string),
+                        txn_snapshot=output["txn_snapshot"],
+                        executor_id=(
+                            output["executor_id"] if output["executor_id"] else None
+                        ),
+                    )
+                )
+        except DBAPIError as dbapi_error:
+            dbos_logger.warning(
+                f'error recording transaction error: {output["error"]}; dbapi_error: {dbapi_error}'
+            )
+            raise dbapi_error
+
     @staticmethod
     def check_transaction_execution(
         session: Session, workflow_uuid: str, function_id: int

dbos/_core.py

@@ -627,7 +627,7 @@ def invoke_tx(*args: Any, **kwargs: Any) -> Any:
                                 assert (
                                     ctx.sql_session is not None
                                 ), "Cannot find a database connection"
-                                ApplicationDatabase.record_transaction_output(
+                                dbos._app_db.record_transaction_output(
                                     ctx.sql_session, txn_output
                                 )
                                 break

dbos/_db_wizard.py

@@ -7,7 +7,7 @@
 import typer
 import yaml
 from rich import print
-from sqlalchemy import URL, create_engine, text
+from sqlalchemy import URL, Engine, create_engine, text
 
 if TYPE_CHECKING:
     from ._dbos_config import ConfigFile
@@ -171,28 +171,46 @@ def _check_docker_installed() -> bool:
 
 
 def _check_db_connectivity(config: "ConfigFile") -> Optional[Exception]:
-    postgres_db_url = URL.create(
-        "postgresql+psycopg",
-        username=config["database"]["username"],
-        password=config["database"]["password"],
-        host=config["database"]["hostname"],
-        port=config["database"]["port"],
-        database="postgres",
-        query={"connect_timeout": "1"},
-    )
-    postgres_db_engine = create_engine(postgres_db_url)
-    try:
-        with postgres_db_engine.connect() as conn:
-            val = conn.execute(text("SELECT 1")).scalar()
-            if val != 1:
-                dbos_logger.error(
-                    f"Unexpected value returned from database: expected 1, received {val}"
-                )
-                return Exception()
-    except Exception as e:
-        return e
-    finally:
-        postgres_db_engine.dispose()
+    database_type = config["database"].get("type", "postgresql")
+
+    engine: Optional[Engine] = None
+    if "postgresql" == database_type:
+        postgres_db_url = URL.create(
+            "postgresql+psycopg",
+            username=config["database"]["username"],
+            password=config["database"]["password"],
+            host=config["database"]["hostname"],
+            port=config["database"]["port"],
+            database="postgres",
+            query={"connect_timeout": "1"},
+        )
+        engine = create_engine(postgres_db_url)
+    elif "mysql" == database_type:
+        db_url_args = {
+            "drivername": "mysql+pymysql",
+            "username": config["database"]["username"],
+            "password": config["database"]["password"],
+            "host": config["database"]["hostname"],
+            "port": config["database"]["port"],
+        }
+        mysql_db_url = URL.create(**db_url_args)
+        engine = create_engine(mysql_db_url)
+
+    if engine:
+        try:
+            with engine.connect() as conn:
+                val = conn.execute(text("SELECT 1")).scalar()
+                if val != 1:
+                    dbos_logger.error(
+                        f"Unexpected value returned from database: expected 1, received {val}"
+                    )
+                    return Exception()
+        except Exception as e:
+            return e
+        finally:
+            engine.dispose()
+    else:
+        return Exception(f"Could not create engine for {database_type} database")
 
     return None
 

dbos/_dbos.py

@@ -417,12 +417,13 @@ def _launch(self) -> None:
             self._executor.submit(startup_recovery_thread, self, workflow_ids)
 
             # Listen to notifications
-            notification_listener_thread = threading.Thread(
-                target=self._sys_db._notification_listener,
-                daemon=True,
-            )
-            notification_listener_thread.start()
-            self._background_threads.append(notification_listener_thread)
+            if self._sys_db.is_notification_listener_enabled():
+                notification_listener_thread = threading.Thread(
+                    target=self._sys_db._notification_listener,
+                    daemon=True,
+                )
+                notification_listener_thread.start()
+                self._background_threads.append(notification_listener_thread)
 
             # Start flush workflow buffers thread
             flush_workflow_buffers_thread = threading.Thread(

dbos/_dbos_config.py

@@ -23,6 +23,7 @@ class RuntimeConfig(TypedDict, total=False):
 
 
 class DatabaseConfig(TypedDict, total=False):
+    type: str
     hostname: str
     port: int
     username: str

dbos/_migrations/mysql/README

@@ -0,0 +1 @@
+Generic single-database configuration.