diff --git a/alembic_data/__init__.py b/alembic_data/__init__.py new file mode 100644 index 0000000..d706f88 --- /dev/null +++ b/alembic_data/__init__.py @@ -0,0 +1,85 @@ +"""Alembic-Data is an extension for alembic that creates automatic data migrations. + +Data migrations (in contrast to schema migrations) are performed on the contents of database tables. They allow you to +create, update, and delete rows inside the database. Alembic-Data provides a framework to generate these migrations +automatically. + +Using Alembic-Data is very simple. You can use the `register_object` function to register an object in the database. +Alembic data will make sure that the next migration creates the registered object. This is normally only useful for +tables that do not contain any user data (as the migrations would have to be generated depending on the actual data). +Instead the typical usecase is to combine the `register_object` function with the `@managed_model` decorator to tell +Alembic-Data that a model will be managed exclusively programmatically. + +The `@managed_model` decorator is used on your declarative model classes. For any `@managed_model` Alembic-Data will +created your `register_object` and will additionally delete any rows from that database that do not belong to any +registered object. +""" + +from typing import List, Any, Dict + +from alembic.autogenerate import comparators +from alembic.autogenerate.api import AutogenContext +from alembic.operations.ops import UpgradeOps +from sqlalchemy.ext.declarative import DeclarativeMeta + +from .ops import * + +__all__ = ["managed_model", "register_object"] + + +def managed_model(cls: DeclarativeMeta) -> DeclarativeMeta: + """Marks this model as a managed model. + + Class decorator to be used on declarative SQLAlchemy models. Using this decorator causes alembic-data to assume that + objects for this model will only be created through code and the `register_object` function. + + :param cls: The class to be decorated. Must be a SQLAlchemy declarative class. + :return: The decorated class, unchanged. + """ + cls.metadata.info.setdefault("managed models", set()).add(cls) # type: ignore + return cls + + +def register_object(cls: DeclarativeMeta, **kwargs) -> None: + """ + Registers an object to be created during a database migration. + + :param cls: The class of the declarative model class. This is used to determine which table to operate on. + :param kwargs: Arguments used to create the object. This must uniquely identify the object as it is also used to + determine whether the object already exists. + """ + cls.metadata.info.setdefault("objects", dict()).setdefault(cls, list()).append(kwargs) # type: ignore + + +@comparators.dispatch_for("schema") +def compare_objects(autogen_context: AutogenContext, upgrade_ops: UpgradeOps, schemas: List[str]) -> None: + """ + Generates appropriate data migrations for the `autogen_context` and adds them to `upgrade_ops`. + """ + managed_models = autogen_context.metadata.info.get('managed models', set()) + objects: Dict[DeclarativeMeta, Any] = autogen_context.metadata.info.get('objects', dict()) + remaining_objects = {cls: set(cls.query.all()) # type: ignore + for cls in objects.keys() + if cls in managed_models + and autogen_context.dialect.has_table(autogen_context.connection.engine, + cls.__table__.name)} # type: ignore + + # Add new objects + for cls, datas in objects.items(): + for data in datas: + existing_object = (cls.query.filter_by(**data).first() # type: ignore + if autogen_context.dialect.has_table(autogen_context.connection.engine, + cls.__table__.name) # type: ignore + else None) + if existing_object: + remaining_objects.get(cls, set()).discard(existing_object) + else: + print("Adding " + cls.__name__ + ": " + str(data)) + upgrade_ops.ops.append(InsertRowOp(cls.__table__.name, **data)) # type: ignore + + # Remove old objects + for cls, object_set in remaining_objects.items(): + for o in object_set: + values = {key: getattr(o, key) for key in o.__mapper__.column_attrs.keys()} + print("Deleting " + cls.__name__ + ": " + str(values)) + upgrade_ops.ops.append(DeleteRowOp(cls.__table__.name, **values)) # type: ignore diff --git a/alembic_data/ops/__init__.py b/alembic_data/ops/__init__.py new file mode 100644 index 0000000..e906059 --- /dev/null +++ b/alembic_data/ops/__init__.py @@ -0,0 +1,5 @@ +"""Operations used by Alembic-Data to create, update and delete objects.""" + +from .delete_row import * +from .insert_row import * +from .update_row import * diff --git a/alembic_data/ops/delete_row.py b/alembic_data/ops/delete_row.py new file mode 100644 index 0000000..e7c0868 --- /dev/null +++ b/alembic_data/ops/delete_row.py @@ -0,0 +1,54 @@ +from typing import Any, Mapping, List + +import sqlalchemy +from alembic.autogenerate import renderers +from alembic.autogenerate.api import AutogenContext +from alembic.operations import Operations, MigrateOperation +from sqlalchemy import and_ + +__all__ = ["DeleteRowOp"] + +from sqlalchemy.sql.elements import ColumnClause + + +@Operations.register_operation("delete_row") +class DeleteRowOp(MigrateOperation): + """This `MigrationOperation` deletes a rows from a database table.""" + + def __init__(self, table_name: str, **kwargs: Mapping[str, Any]) -> None: + """Creates a new `DeleteRowOp`. + + :param table_name: The name of the table to delete the row from. + :param kwargs: The values used to identify the exact row. Keys must correspond to column names in the table and + values to their values. + """ + self.table_name = table_name + self.values = kwargs + + @classmethod + def delete_row(cls, operations: Operations, table_name: str, **kwargs: Mapping[str, Any]): + """Factory method for `DeleteRowOp`.""" + op = DeleteRowOp(table_name, **kwargs) + return operations.invoke(op) + + def reverse(self) -> MigrateOperation: + """Reverses the insert row operation by inserting the respective row.""" + from .insert_row import InsertRowOp + return InsertRowOp(self.table_name, **self.values) + + +@Operations.implementation_for(DeleteRowOp) +def delete_row(operations, operation: DeleteRowOp) -> None: + columns: List[ColumnClause] = [sqlalchemy.column(key) for key in operation.values.keys()] + table = sqlalchemy.table(operation.table_name, *columns) + where_clauses = [table.columns[key] == value for (key, value) in operation.values.items()] + sql = table.delete().where(and_(*where_clauses)) + operations.execute(sql) + + +@renderers.dispatch_for(DeleteRowOp) +def render_delete_row(autogen_context: AutogenContext, op: DeleteRowOp) -> str: + return "op.delete_row(%r, **%r)" % ( + op.table_name, + op.values + ) diff --git a/alembic_data/ops/insert_row.py b/alembic_data/ops/insert_row.py new file mode 100644 index 0000000..750ed12 --- /dev/null +++ b/alembic_data/ops/insert_row.py @@ -0,0 +1,51 @@ +from typing import Any, Mapping, List + +import sqlalchemy +from alembic.autogenerate import renderers +from alembic.autogenerate.api import AutogenContext +from alembic.operations import Operations, MigrateOperation + +__all__ = ["InsertRowOp"] + +from sqlalchemy.sql.elements import ColumnClause + + +@Operations.register_operation("insert_row") +class InsertRowOp(MigrateOperation): + """This `MigrationOperation` adds rows to a database table.""" + + def __init__(self, table_name: str, **kwargs: Mapping[str, Any]) -> None: + """Creates a new `InsertRowOp`. + + :param table_name: The name of the table to insert the row. + :param kwargs: The values to insert into the table. Keys must correspond to column names in the table. + """ + self.table_name = table_name + self.values = kwargs + + @classmethod + def insert_row(cls, operations: Operations, table_name: str, **kwargs: Mapping[str, Any]): + """Factory method for `InsertRowOp`.""" + op = InsertRowOp(table_name, **kwargs) + return operations.invoke(op) + + def reverse(self) -> MigrateOperation: + """Reverses the insert row operation by deleting the respective row.""" + from .delete_row import DeleteRowOp + return DeleteRowOp(self.table_name, **self.values) + + +@Operations.implementation_for(InsertRowOp) +def insert_row(operations: Operations, operation: InsertRowOp) -> None: + columns: List[ColumnClause] = [sqlalchemy.column(key) for key in operation.values.keys()] + table = sqlalchemy.table(operation.table_name, *columns) + sql = table.insert().values(**operation.values) + operations.execute(sql) + + +@renderers.dispatch_for(InsertRowOp) +def render_insert_row(autogen_context: AutogenContext, op: InsertRowOp) -> str: + return "op.insert_row(%r, **%r)" % ( + op.table_name, + op.values + ) diff --git a/alembic_data/ops/update_row.py b/alembic_data/ops/update_row.py new file mode 100644 index 0000000..fc20a7b --- /dev/null +++ b/alembic_data/ops/update_row.py @@ -0,0 +1,57 @@ +from typing import Any, Mapping, List + +import sqlalchemy +from alembic.autogenerate import renderers +from alembic.autogenerate.api import AutogenContext +from alembic.operations import Operations, MigrateOperation + +__all__ = ["UpdateRowOp"] + +from sqlalchemy.sql.elements import ColumnClause + + +@Operations.register_operation("update_row") +class UpdateRowOp(MigrateOperation): + """This `MigrationOperation` updates rows in a database table.""" + + def __init__(self, table_name: str, new_values: Mapping[str, Any], old_values: Mapping[str, Any]) -> None: + """Creates a new `InsertRowOp`. + + :param table_name: The name of the table to update values in. + :param new_values: The new values to update in the table. Keys must correspond to column names in the table and + values to their values. + :param old_values: The old values used to identify the row that should be updated. + """ + self.table_name = table_name + self.new_values = new_values + self.old_values = old_values + + @classmethod + def update_row(cls, operations: Operations, table_name: str, new_values: Mapping[str, Any], + old_values: Mapping[str, Any]): + """Factory method for `UpdateRowOp`.""" + op = UpdateRowOp(table_name, new_values, old_values) + return operations.invoke(op) + + def reverse(self) -> MigrateOperation: + """Reverses the update operation by reverting back to the old values.""" + return UpdateRowOp(self.table_name, self.old_values, self.new_values) + + +@Operations.implementation_for(UpdateRowOp) +def update_row(operations: Operations, operation: UpdateRowOp) -> None: + columns: List[ColumnClause] = [sqlalchemy.column(key) for key in operation.values.keys()] + table = sqlalchemy.table(operation.table_name, *columns) + where_clauses = [table.columns[key] == value for (key, value) in operation.old_values.items()] + operations.execute( + table.update().where(*where_clauses).values(**operation.new_values) + ) + + +@renderers.dispatch_for(UpdateRowOp) +def render_update_row(autogen_context: AutogenContext, op: UpdateRowOp) -> str: + return "op.update_row(%r, **%r, **%r)" % ( + op.table_name, + op.new_values, + op.old_values + ) diff --git a/alembic_data/py.typed b/alembic_data/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..9c58808 --- /dev/null +++ b/setup.py @@ -0,0 +1,47 @@ +""" +Alembic-Data +------------ + +Alembic-Data is an extension for alembic that allows the automatic creation of database content (in addition to schema). +It adds +""" +from setuptools import setup + +setup( + name='Alembic-Data', + version='0.9', + url='https://gitlab.com/Codello/alembic-data/', + license='MIT', + author='Kim Wittenburg', + author_email='codello@wittenburg.kim', + description='Data Migrations for Alembic.', + long_description=__doc__, + packages=['alembic_data'], + package_data={"alembic_data": ["py.typed"]}, + zip_safe=False, + include_package_data=True, + platforms='any', + install_requires=[ + 'alembic' + ], + extras_require={ + 'types': [ + 'sqlalchemy-stubs' + ] + }, + test_requirements=[ + 'pytest', + 'pytest-pep8', + 'pytest-cov', + ], + classifiers=[ + 'Development Status :: 3 - Alpha', + 'Environment :: Plugins', + 'Intended Audience :: Developers', + 'License :: OSI Approved :: MIT License', + 'Operating System :: OS Independent', + 'Programming Language :: Python', + 'Topic :: Database', + 'Topic :: Software Development :: Libraries :: Python Modules' + ] +)