-
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.
* multiple slackbot support * app_id + tenant_id key * removed kv store stuff * fixed up mypy and migration * got frontend working for multiple slack bots * some frontend stuff * alembic fix * might be valid * refactor dun * alembic stuff * temp frontend stuff * alembic stuff * maybe fixed alembic * maybe dis fix * im getting mad * api names changed * tested * almost done * done * routing nonsense * done! * done!! * fr done * doneski * fix alembic migration * getting mad again * PLEASE IM BEGGING YOU
- Loading branch information
1 parent
b712877
commit 9209fc8
Showing
59 changed files
with
2,422 additions
and
1,026 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
280 changes: 280 additions & 0 deletions
280
backend/alembic/versions/4ee1287bd26a_add_multiple_slack_bot_support.py
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,280 @@ | ||
| """add_multiple_slack_bot_support | ||
| Revision ID: 4ee1287bd26a | ||
| Revises: 47e5bef3a1d7 | ||
| Create Date: 2024-11-06 13:15:53.302644 | ||
| """ | ||
| import logging | ||
| from typing import cast | ||
| from alembic import op | ||
| import sqlalchemy as sa | ||
| from sqlalchemy.orm import Session | ||
| from danswer.key_value_store.factory import get_kv_store | ||
| from danswer.db.models import SlackBot | ||
| from sqlalchemy.dialects import postgresql | ||
|
|
||
| # revision identifiers, used by Alembic. | ||
| revision = "4ee1287bd26a" | ||
| down_revision = "47e5bef3a1d7" | ||
| branch_labels: None = None | ||
| depends_on: None = None | ||
|
|
||
| # Configure logging | ||
| logger = logging.getLogger("alembic.runtime.migration") | ||
| logger.setLevel(logging.INFO) | ||
|
|
||
|
|
||
| def upgrade() -> None: | ||
| logger.info(f"{revision}: create_table: slack_bot") | ||
| # Create new slack_bot table | ||
| op.create_table( | ||
| "slack_bot", | ||
| sa.Column("id", sa.Integer(), nullable=False), | ||
| sa.Column("name", sa.String(), nullable=False), | ||
| sa.Column("enabled", sa.Boolean(), nullable=False, server_default="true"), | ||
| sa.Column("bot_token", sa.LargeBinary(), nullable=False), | ||
| sa.Column("app_token", sa.LargeBinary(), nullable=False), | ||
| sa.PrimaryKeyConstraint("id"), | ||
| sa.UniqueConstraint("bot_token"), | ||
| sa.UniqueConstraint("app_token"), | ||
| ) | ||
|
|
||
| # # Create new slack_channel_config table | ||
| op.create_table( | ||
| "slack_channel_config", | ||
| sa.Column("id", sa.Integer(), nullable=False), | ||
| sa.Column("slack_bot_id", sa.Integer(), nullable=True), | ||
| sa.Column("persona_id", sa.Integer(), nullable=True), | ||
| sa.Column("channel_config", postgresql.JSONB(), nullable=False), | ||
| sa.Column("response_type", sa.String(), nullable=False), | ||
| sa.Column( | ||
| "enable_auto_filters", sa.Boolean(), nullable=False, server_default="false" | ||
| ), | ||
| sa.ForeignKeyConstraint( | ||
| ["slack_bot_id"], | ||
| ["slack_bot.id"], | ||
| ), | ||
| sa.ForeignKeyConstraint( | ||
| ["persona_id"], | ||
| ["persona.id"], | ||
| ), | ||
| sa.PrimaryKeyConstraint("id"), | ||
| ) | ||
|
|
||
| # Handle existing Slack bot tokens first | ||
| logger.info(f"{revision}: Checking for existing Slack bot.") | ||
| bot_token = None | ||
| app_token = None | ||
| first_row_id = None | ||
|
|
||
| try: | ||
| tokens = cast(dict, get_kv_store().load("slack_bot_tokens_config_key")) | ||
| except Exception: | ||
| logger.warning("No existing Slack bot tokens found.") | ||
| tokens = {} | ||
|
|
||
| bot_token = tokens.get("bot_token") | ||
| app_token = tokens.get("app_token") | ||
|
|
||
| if bot_token and app_token: | ||
| logger.info(f"{revision}: Found bot and app tokens.") | ||
|
|
||
| session = Session(bind=op.get_bind()) | ||
| new_slack_bot = SlackBot( | ||
| name="Slack Bot (Migrated)", | ||
| enabled=True, | ||
| bot_token=bot_token, | ||
| app_token=app_token, | ||
| ) | ||
| session.add(new_slack_bot) | ||
| session.commit() | ||
| first_row_id = new_slack_bot.id | ||
|
|
||
| # Create a default bot if none exists | ||
| # This is in case there are no slack tokens but there are channels configured | ||
| op.execute( | ||
| sa.text( | ||
| """ | ||
| INSERT INTO slack_bot (name, enabled, bot_token, app_token) | ||
| SELECT 'Default Bot', true, '', '' | ||
| WHERE NOT EXISTS (SELECT 1 FROM slack_bot) | ||
| RETURNING id; | ||
| """ | ||
| ) | ||
| ) | ||
|
|
||
| # Get the bot ID to use (either from existing migration or newly created) | ||
| bot_id_query = sa.text( | ||
| """ | ||
| SELECT COALESCE( | ||
| :first_row_id, | ||
| (SELECT id FROM slack_bot ORDER BY id ASC LIMIT 1) | ||
| ) as bot_id; | ||
| """ | ||
| ) | ||
| result = op.get_bind().execute(bot_id_query, {"first_row_id": first_row_id}) | ||
| bot_id = result.scalar() | ||
|
|
||
| # CTE (Common Table Expression) that transforms the old slack_bot_config table data | ||
| # This splits up the channel_names into their own rows | ||
| channel_names_cte = """ | ||
| WITH channel_names AS ( | ||
| SELECT | ||
| sbc.id as config_id, | ||
| sbc.persona_id, | ||
| sbc.response_type, | ||
| sbc.enable_auto_filters, | ||
| jsonb_array_elements_text(sbc.channel_config->'channel_names') as channel_name, | ||
| sbc.channel_config->>'respond_tag_only' as respond_tag_only, | ||
| sbc.channel_config->>'respond_to_bots' as respond_to_bots, | ||
| sbc.channel_config->'respond_member_group_list' as respond_member_group_list, | ||
| sbc.channel_config->'answer_filters' as answer_filters, | ||
| sbc.channel_config->'follow_up_tags' as follow_up_tags | ||
| FROM slack_bot_config sbc | ||
| ) | ||
| """ | ||
|
|
||
| # Insert the channel names into the new slack_channel_config table | ||
| insert_statement = """ | ||
| INSERT INTO slack_channel_config ( | ||
| slack_bot_id, | ||
| persona_id, | ||
| channel_config, | ||
| response_type, | ||
| enable_auto_filters | ||
| ) | ||
| SELECT | ||
| :bot_id, | ||
| channel_name.persona_id, | ||
| jsonb_build_object( | ||
| 'channel_name', channel_name.channel_name, | ||
| 'respond_tag_only', | ||
| COALESCE((channel_name.respond_tag_only)::boolean, false), | ||
| 'respond_to_bots', | ||
| COALESCE((channel_name.respond_to_bots)::boolean, false), | ||
| 'respond_member_group_list', | ||
| COALESCE(channel_name.respond_member_group_list, '[]'::jsonb), | ||
| 'answer_filters', | ||
| COALESCE(channel_name.answer_filters, '[]'::jsonb), | ||
| 'follow_up_tags', | ||
| COALESCE(channel_name.follow_up_tags, '[]'::jsonb) | ||
| ), | ||
| channel_name.response_type, | ||
| channel_name.enable_auto_filters | ||
| FROM channel_names channel_name; | ||
| """ | ||
|
|
||
| op.execute(sa.text(channel_names_cte + insert_statement).bindparams(bot_id=bot_id)) | ||
|
|
||
| # Clean up old tokens if they existed | ||
| try: | ||
| if bot_token and app_token: | ||
| logger.info(f"{revision}: Removing old bot and app tokens.") | ||
| get_kv_store().delete("slack_bot_tokens_config_key") | ||
| except Exception: | ||
| logger.warning("tried to delete tokens in dynamic config but failed") | ||
| # Rename the table | ||
| op.rename_table( | ||
| "slack_bot_config__standard_answer_category", | ||
| "slack_channel_config__standard_answer_category", | ||
| ) | ||
|
|
||
| # Rename the column | ||
| op.alter_column( | ||
| "slack_channel_config__standard_answer_category", | ||
| "slack_bot_config_id", | ||
| new_column_name="slack_channel_config_id", | ||
| ) | ||
|
|
||
| # Drop the table with CASCADE to handle dependent objects | ||
| op.execute("DROP TABLE slack_bot_config CASCADE") | ||
|
|
||
| logger.info(f"{revision}: Migration complete.") | ||
|
|
||
|
|
||
| def downgrade() -> None: | ||
| # Recreate the old slack_bot_config table | ||
| op.create_table( | ||
| "slack_bot_config", | ||
| sa.Column("id", sa.Integer(), nullable=False), | ||
| sa.Column("persona_id", sa.Integer(), nullable=True), | ||
| sa.Column("channel_config", postgresql.JSONB(), nullable=False), | ||
| sa.Column("response_type", sa.String(), nullable=False), | ||
| sa.Column("enable_auto_filters", sa.Boolean(), nullable=False), | ||
| sa.ForeignKeyConstraint( | ||
| ["persona_id"], | ||
| ["persona.id"], | ||
| ), | ||
| sa.PrimaryKeyConstraint("id"), | ||
| ) | ||
|
|
||
| # Migrate data back to the old format | ||
| # Group by persona_id to combine channel names back into arrays | ||
| op.execute( | ||
| sa.text( | ||
| """ | ||
| INSERT INTO slack_bot_config ( | ||
| persona_id, | ||
| channel_config, | ||
| response_type, | ||
| enable_auto_filters | ||
| ) | ||
| SELECT DISTINCT ON (persona_id) | ||
| persona_id, | ||
| jsonb_build_object( | ||
| 'channel_names', ( | ||
| SELECT jsonb_agg(c.channel_config->>'channel_name') | ||
| FROM slack_channel_config c | ||
| WHERE c.persona_id = scc.persona_id | ||
| ), | ||
| 'respond_tag_only', (channel_config->>'respond_tag_only')::boolean, | ||
| 'respond_to_bots', (channel_config->>'respond_to_bots')::boolean, | ||
| 'respond_member_group_list', channel_config->'respond_member_group_list', | ||
| 'answer_filters', channel_config->'answer_filters', | ||
| 'follow_up_tags', channel_config->'follow_up_tags' | ||
| ), | ||
| response_type, | ||
| enable_auto_filters | ||
| FROM slack_channel_config scc | ||
| WHERE persona_id IS NOT NULL; | ||
| """ | ||
| ) | ||
| ) | ||
|
|
||
| # Rename the table back | ||
| op.rename_table( | ||
| "slack_channel_config__standard_answer_category", | ||
| "slack_bot_config__standard_answer_category", | ||
| ) | ||
|
|
||
| # Rename the column back | ||
| op.alter_column( | ||
| "slack_bot_config__standard_answer_category", | ||
| "slack_channel_config_id", | ||
| new_column_name="slack_bot_config_id", | ||
| ) | ||
|
|
||
| # Try to save the first bot's tokens back to KV store | ||
| try: | ||
| first_bot = ( | ||
| op.get_bind() | ||
| .execute( | ||
| sa.text( | ||
| "SELECT bot_token, app_token FROM slack_bot ORDER BY id LIMIT 1" | ||
| ) | ||
| ) | ||
| .first() | ||
| ) | ||
| if first_bot and first_bot.bot_token and first_bot.app_token: | ||
| tokens = { | ||
| "bot_token": first_bot.bot_token, | ||
| "app_token": first_bot.app_token, | ||
| } | ||
| get_kv_store().store("slack_bot_tokens_config_key", tokens) | ||
| except Exception: | ||
| logger.warning("Failed to save tokens back to KV store") | ||
|
|
||
| # Drop the new tables in reverse order | ||
| op.drop_table("slack_channel_config") | ||
| op.drop_table("slack_bot") |
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
2 changes: 1 addition & 1 deletion
2
backend/alembic/versions/9cf5c00f72fe_add_creator_to_cc_pair.py
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
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
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
Oops, something went wrong.