From 9d64a4aa9d8338a7d2415c09c32fb646c9111c4b Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Fri, 22 Dec 2023 21:15:08 +0100 Subject: [PATCH] Add migration script --- .../7b651bf6822e_track_secrets_in_db.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 src/zenml/zen_stores/migrations/versions/7b651bf6822e_track_secrets_in_db.py diff --git a/src/zenml/zen_stores/migrations/versions/7b651bf6822e_track_secrets_in_db.py b/src/zenml/zen_stores/migrations/versions/7b651bf6822e_track_secrets_in_db.py new file mode 100644 index 00000000000..73ba064fe36 --- /dev/null +++ b/src/zenml/zen_stores/migrations/versions/7b651bf6822e_track_secrets_in_db.py @@ -0,0 +1,58 @@ +"""track secrets in db [7b651bf6822e]. + +Revision ID: 7b651bf6822e +Revises: 0.53.1 +Create Date: 2023-12-22 20:56:18.131906 + +""" +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. + + +revision = "7b651bf6822e" +down_revision = "0.53.1" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + """Upgrade database schema and/or data, creating a new revision.""" + from zenml.config.global_config import GlobalConfiguration + from zenml.enums import SecretsStoreType + + zen_store = GlobalConfiguration().zen_store + + if zen_store.secrets_store.type == SecretsStoreType.SQL: + return + + # Transfer secrets from the external secrets store to the db + secrets_store = zen_store.secrets_store + + external_secrets = secrets_store.list_secrets() + + conn = op.get_bind() + meta = sa.MetaData(bind=op.get_bind()) + meta.reflect(only=("secret",)) + secrets = sa.Table("artifact", meta) + for secret in external_secrets: + conn.execute( + secrets.insert().values( + id=secret.id, + created=secret.created, + updated=secret.updated, + name=secret.name, + scope=secret.scope, + values=None, + workspace_id=secret.workspace.id, + user_id=secret.user.id if secret.user else None, + ) + ) + + +def downgrade() -> None: + """Downgrade database schema and/or data back to the previous revision.""" + # ### commands auto generated by Alembic - please adjust! ### + pass + # ### end Alembic commands ###