MySQL schema change to support fractional time (#984)

# Pull Request

## Title

MySQL schema change to support fractional time

______________________________________________________________________

## Description

- Changes the storage schema to support a variant for MySQL that can
handle fractional time.
- Adds a custom column comparator to be able to pick that up by default
in the future.
- Adjusts tests to check for fractional time support.

______________________________________________________________________

## Type of Change

- 🛠️ Bug fix
- ✨ New feature
- 🔄 Refactor
- 🧪 Tests

______________________________________________________________________

## Testing

- A lot of manual testing with docker.
- Additional CI tests
- Existing CI tests

______________________________________________________________________

## Additional Notes (optional)

Part of a series of changes, should be merged after #983 

______________________________________________________________________

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bpkroth

mlos_bench/mlos_bench/storage/sql/alembic.ini

@@ -63,7 +63,10 @@ version_path_separator = os  # Use os.pathsep. Default configuration used for ne
 # output_encoding = utf-8
 
 # See README.md for details.
+# Uncomment one of these:
 sqlalchemy.url = sqlite:///mlos_bench.sqlite
+#sqlalchemy.url = mysql+mysqlconnector://root:password@localhost:3306/mlos_bench
+#sqlalchemy.url = postgresql+psycopg2://root:password@localhost:5432/mlos_bench
 
 
 [post_write_hooks]
@@ -72,10 +75,10 @@ sqlalchemy.url = sqlite:///mlos_bench.sqlite
 # detail and examples
 
 # format using "black" - use the console_scripts runner, against the "black" entrypoint
-# hooks = black
-# black.type = console_scripts
-# black.entrypoint = black
-# black.options = -l 79 REVISION_SCRIPT_FILENAME
+hooks = black
+black.type = console_scripts
+black.entrypoint = black
+black.options = REVISION_SCRIPT_FILENAME
 
 # lint with attempts to fix using "ruff" - use the exec runner, execute a binary
 # hooks = ruff

mlos_bench/mlos_bench/storage/sql/alembic/README.md

@@ -4,49 +4,144 @@ This document contains some notes on how to use [`alembic`](https://alembic.sqla
 
 ## Overview
 
-1. Create a blank `mlos_bench.sqlite` database file in the [`mlos_bench/storage/sql`](../) directory with the current schema using the following command:
+1. Create a blank database instance in the [`mlos_bench/storage/sql`](../) directory with the current schema using the following command:
 
-   ```sh
-   cd mlos_bench/storage/sql
-   rm mlos_bench.sqlite
-   mlos_bench --storage storage/sqlite.jsonc --create-update-storage-schema-only
-   ```
+   This allows `alembic` to automatically generate a migration script from the current schema.
+
+   > NOTE: If your schema changes target a particular backend engine (e.g., using `with_variant`) you will need to use an engine with that config for this step.
+   > \
+   > In the remainder of this document we should some examples for different DB types.
+   > Pick the one you're targeting and stick with it thru the example.
+   > You may need to repeat the process several times to test all of them.
+   >
+   > - [ ] TODO: Add scripts to automatically do this for several different backend engines all at once.
+
+   For instance:
+
+   1. Start a temporary server either as a local file or in a docker instance
+
+      ```sh
+      # sqlite
+      cd mlos_bench/storage/sql
+      rm -f mlos_bench.sqlite
+      ```
+
+      ```sh
+      # mysql
+      docker run -it --rm --name mysql-alembic --env MYSQL_ROOT_PASSWORD=password --env MYSQL_DATABASE=mlos_bench -p 3306:3306 mysql:latest
+      ```
+
+      ```sh
+      # postgres
+      docker run -it --rm --name postgres-alembic --env POSTGRES_PASSWORD=password --env POSTGRES_DB=mlos_bench -p 5432:5432 postgres:latest
+      ```
+
+   1. Adjust the `sqlalchemy.url` in the [`alembic.ini`](../alembic.ini) file.
 
-   > This allows `alembic` to automatically generate a migration script from the current schema.
+      ```ini
+      # Uncomment one of these.
+      # See README.md for details.
 
-1. Adjust the [`mlos_bench/storage/sql/schema.py`](../schema.py) file to reflect the new desired schema.
+      #sqlalchemy.url = sqlite:///mlos_bench.sqlite
+      sqlalchemy.url = mysql+pymysql://root:password@localhost:3306/mlos_bench
+      #sqlalchemy.url = postgresql+psycopg2://root:password@localhost:5432/mlos_bench
+      ```
 
+   1. Prime the DB schema
+
+      > Note: you may want to `git checkout main` first to make sure you're using the current schema here.
+
+      ```sh
+      # sqlite
+      mlos_bench --storage storage/sqlite.jsonc --create-update-storage-schema-only --password=password
+      ```
+
+      ```sh
+      # mysql
+      mlos_bench --storage storage/mysql.jsonc --create-update-storage-schema-only --password=password
+      ```
+
+      ```sh
+      # postgres
+      mlos_bench --storage storage/postgresql.jsonc --create-update-storage-schema-only --password=password
+      ```
+
+1. Now, adjust the [`mlos_bench/storage/sql/schema.py`](../schema.py) file to reflect the new desired schema.
+
+   > Don't forget to do this on a new branch.
+   > \
    > Keep each change small and atomic.
+   > \
    > For example, if you want to add a new column, do that in one change.
    > If you want to rename a column, do that in another change.
 
 1. Generate a new migration script with the following command:
 
    ```sh
-   alembic revision --autogenerate -m "Descriptive text about the change."
+   alembic revision --autogenerate -m "CHANGEME: Descriptive text about the change."
    ```
 
 1. Review the generated migration script in the [`mlos_bench/storage/sql/alembic/versions`](./versions/) directory.
 
 1. Verify that the migration script works by running the following command:
 
    ```sh
+   # sqlite
    mlos_bench --storage storage/sqlite.jsonc --create-update-storage-schema-only
    ```
 
+   ```sh
+   # mysql:
+   mlos_bench --storage storage/mysql.jsonc --create-update-storage-schema-only --password=password
+   ```
+
+   ```sh
+   # postgres:
+   mlos_bench --storage storage/postgresql.jsonc --create-update-storage-schema-only --password=password
+   ```
+
    > Normally this would be done with `alembic upgrade head`, but this command is convenient to ensure if will work with the `mlos_bench` command line interface as well.
 
    Examine the results using something like:
 
    ```sh
+   # For sqlite:
    sqlite3 mlos_bench.sqlite .schema
    sqlite3 mlos_bench.sqlite "SELECT * FROM alembic_version;"
    ```
 
+   ```sh
+   # For mysql:
+   mysql --user root --password=password --host localhost --protocol tcp --database mlos_bench -e "SHOW TABLES; SELECT * FROM alembic_version;"
+   ```
+
+   ```sh
+   # For postgres:
+   PGPASSWORD=password psql -h localhost -p 5432 -U postgres mlos_bench -c "SELECT table_name FROM information_schema.tables WHERE table_schema='public' and table_catalog='mlos_bench'; SELECT * FROM alembic_version;"
+   ```
+
+   > Use different CLI clients for targeting other engines.
+
 1. If the migration script works, commit the changes to the [`mlos_bench/storage/sql/schema.py`](../schema.py) and [`mlos_bench/storage/sql/alembic/versions`](./versions/) files.
 
    > Be sure to update the latest version in the [`test_storage_schemas.py`](../../../tests/storage/sql/test_storage_schemas.py) file as well.
 
+1. Cleanup any server instances you started.
+
+   For instance:
+
+   ```sh
+   rm mlos_bench/storage/sql/mlos_bench.sqlite
+   ```
+
+   ```sh
+   docker kill mysql-alembic
+   ```
+
+   ```sh
+   docker kill postgres-alembic
+   ```
+
 1. Merge that to the `main` branch.
 
 1. Might be good to cut a new `mlos_bench` release at this point as well.

mlos_bench/mlos_bench/storage/sql/alembic/env.py

@@ -5,11 +5,17 @@
 """Alembic environment script."""
 # pylint: disable=no-member
 
+import logging
 import sys
 from logging.config import fileConfig
 
 from alembic import context
-from sqlalchemy import engine_from_config, pool
+from alembic.migration import MigrationContext
+from sqlalchemy import create_engine, engine_from_config, pool
+from sqlalchemy.dialects import mysql
+from sqlalchemy.schema import Column as SchemaColumn
+from sqlalchemy.sql.schema import Column as MetadataColumn
+from sqlalchemy.types import TypeEngine
 
 from mlos_bench.storage.sql.schema import DbSchema
 
@@ -22,17 +28,137 @@
 # Don't override the mlos_bench or pytest loggers though.
 if config.config_file_name is not None and "alembic" in sys.argv[0]:
     fileConfig(config.config_file_name)
+alembic_logger = logging.getLogger("alembic")
 
 # add your model's MetaData object here
 # for 'autogenerate' support
-target_metadata = DbSchema(engine=None).meta
+# NOTE: We override the alembic.ini file programmatically in storage/sql/schema.py
+# However, the alembic.ini file value is used during alembic CLI operations
+# (e.g., dev ops extending the schema).
+sqlalchemy_url = config.get_main_option("sqlalchemy.url")
+if not sqlalchemy_url:
+    raise ValueError("Missing sqlalchemy.url: schema changes may not be accurate.")
+engine = create_engine(sqlalchemy_url)
+alembic_logger.info("engine.url %s", str(engine.url))
+target_metadata = DbSchema(engine=engine).meta
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
 
 
+def fq_class_name(t: object) -> str:
+    """Return the fully qualified class name of a type."""
+    return t.__module__ + "." + t.__class__.__name__
+
+
+def custom_compare_types(
+    migration_context: MigrationContext,
+    inspected_column: SchemaColumn | None,
+    metadata_column: MetadataColumn,
+    inspected_type: TypeEngine,
+    metadata_type: TypeEngine,
+) -> bool | None:
+    """
+    Custom column type comparator.
+
+    See `Comparing Types
+    <https://alembic.sqlalchemy.org/en/latest/autogenerate.html#comparing-types>`_
+    documentation for more details.
+
+    Notes
+    -----
+    In the case of a MySQL DateTime variant, it makes sure that the floating
+    point accuracy is met.
+
+    Returns
+    -------
+    result : bool | None
+        Returns True if the column specifications don't match the column (i.e.,
+        a change is needed).
+        Returns False when the column specification and column match.
+        Returns None to fallback to the default comparator logic.
+    """
+    metadata_dialect_type = metadata_type.dialect_impl(migration_context.dialect)
+    if alembic_logger.isEnabledFor(logging.DEBUG):
+        alembic_logger.debug(
+            (
+                "Comparing columns: "
+                "inspected_column: [%s] %s and "
+                "metadata_column: [%s (%s)] %s "
+                "inspected_column.__dict__: %s\n"
+                "inspected_column.dialect_options: %s\n"
+                "inspected_column.dialect_kwargs: %s\n"
+                "inspected_type.__dict__: %s\n"
+                "metadata_column.__dict__: %s\n"
+                "metadata_type.__dict__: %s\n"
+                "metadata_dialect_type.__dict__: %s\n"
+            ),
+            fq_class_name(inspected_type),
+            inspected_column,
+            fq_class_name(metadata_type),
+            fq_class_name(metadata_dialect_type),
+            metadata_column,
+            inspected_column.__dict__,
+            dict(inspected_column.dialect_options) if inspected_column is not None else None,
+            dict(inspected_column.dialect_kwargs) if inspected_column is not None else None,
+            inspected_type.__dict__,
+            metadata_column.__dict__,
+            metadata_type.__dict__,
+            metadata_dialect_type.__dict__,
+        )
+
+    # Implement a more detailed DATETIME precision comparison for MySQL.
+    # Note: Currently also handles MariaDB.
+    if migration_context.dialect.name == "mysql":
+        if isinstance(metadata_dialect_type, (mysql.DATETIME, mysql.TIMESTAMP)):
+            if not isinstance(inspected_type, type(metadata_dialect_type)):
+                alembic_logger.info(
+                    "inspected_type %s does not match metadata_dialect_type %s",
+                    fq_class_name(inspected_type),
+                    fq_class_name(metadata_dialect_type),
+                )
+                return True
+            else:
+                assert isinstance(
+                    inspected_type, (mysql.DATETIME, mysql.TIMESTAMP)
+                ), "Expected inspected_type to be a MySQL DATETIME or TIMESTAMP type."
+                if inspected_type.fsp != metadata_dialect_type.fsp:
+                    alembic_logger.info(
+                        "inspected_type.fsp (%s) and metadata_dialect_type.fsp (%s) don't match",
+                        inspected_type.fsp,
+                        metadata_dialect_type.fsp,
+                    )
+                    return True
+
+                if inspected_type.timezone != metadata_dialect_type.timezone:
+                    alembic_logger.info(
+                        (
+                            "inspected_type.timezone (%s) and "
+                            "metadata_dialect_type.timezone (%s) don't match"
+                        ),
+                        inspected_type.timezone,
+                        metadata_dialect_type.timezone,
+                    )
+                    return True
+
+    if alembic_logger.isEnabledFor(logging.DEBUG):
+        alembic_logger.debug(
+            (
+                "Using default compare_type behavior for "
+                "inspected_column: [%s] %s and "
+                "metadata_column: [%s (%s)] %s (see above for details).\n"
+            ),
+            fq_class_name(inspected_type),
+            inspected_column,
+            fq_class_name(metadata_type),
+            fq_class_name(metadata_dialect_type),
+            metadata_column,
+        )
+    return None  # fallback to default comparison behavior
+
+
 def run_migrations_offline() -> None:
     """
     Run migrations in 'offline' mode.
@@ -49,6 +175,7 @@ def run_migrations_offline() -> None:
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
+        compare_type=custom_compare_types,
     )
 
     with context.begin_transaction():
@@ -74,12 +201,20 @@ def run_migrations_online() -> None:
         )
 
         with connectable.connect() as connection:
-            context.configure(connection=connection, target_metadata=target_metadata)
+            context.configure(
+                connection=connection,
+                target_metadata=target_metadata,
+                compare_type=custom_compare_types,
+            )
 
             with context.begin_transaction():
                 context.run_migrations()
     else:
-        context.configure(connection=connectable, target_metadata=target_metadata)
+        context.configure(
+            connection=connectable,
+            target_metadata=target_metadata,
+            compare_type=custom_compare_types,
+        )
 
         with context.begin_transaction():
             context.run_migrations()