diff --git a/backend/alembic/versions/c0aab6edb6dd_delete_workspace.py b/backend/alembic/versions/c0aab6edb6dd_delete_workspace.py new file mode 100644 index 00000000000..0ead6dd907f --- /dev/null +++ b/backend/alembic/versions/c0aab6edb6dd_delete_workspace.py @@ -0,0 +1,87 @@ +"""delete workspace + +Revision ID: c0aab6edb6dd +Revises: 35e518e0ddf4 +Create Date: 2024-12-17 14:37:07.660631 + +""" + +from alembic import op + + +# revision identifiers, used by Alembic. +revision = "c0aab6edb6dd" +down_revision = "35e518e0ddf4" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.execute( + """ + UPDATE connector + SET connector_specific_config = connector_specific_config - 'workspace' + WHERE source = 'SLACK' + """ + ) + + +def downgrade() -> None: + import json + from sqlalchemy import text + from slack_sdk import WebClient + + conn = op.get_bind() + + # Fetch all Slack credentials + creds_result = conn.execute( + text("SELECT id, credential_json FROM credential WHERE source = 'SLACK'") + ) + all_slack_creds = creds_result.fetchall() + if not all_slack_creds: + return + + for cred_row in all_slack_creds: + credential_id, credential_json = cred_row + + credential_json = ( + credential_json.tobytes().decode("utf-8") + if isinstance(credential_json, memoryview) + else credential_json.decode("utf-8") + ) + credential_data = json.loads(credential_json) + slack_bot_token = credential_data.get("slack_bot_token") + if not slack_bot_token: + print( + f"No slack_bot_token found for credential {credential_id}. " + "Your Slack connector will not function until you upgrade and provide a valid token." + ) + continue + + client = WebClient(token=slack_bot_token) + try: + auth_response = client.auth_test() + workspace = auth_response["url"].split("//")[1].split(".")[0] + + # Update only the connectors linked to this credential + # (and which are Slack connectors). + op.execute( + f""" + UPDATE connector AS c + SET connector_specific_config = jsonb_set( + connector_specific_config, + '{{workspace}}', + to_jsonb('{workspace}'::text) + ) + FROM connector_credential_pair AS ccp + WHERE ccp.connector_id = c.id + AND c.source = 'SLACK' + AND ccp.credential_id = {credential_id} + """ + ) + except Exception: + print( + f"We were unable to get the workspace url for your Slack Connector with id {credential_id}." + ) + print("This connector will no longer work until you upgrade.") + continue