-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Old alembic migration + restore workspace
- Loading branch information
Showing
1 changed file
with
87 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |