Merge pull request #855 from Fr1ngg/iex73b-bedolaga/-api

Add user-linked API tokens and issuance flow
This commit is contained in:
Egor
2025-10-08 01:19:44 +03:00
committed by GitHub
11 changed files with 448 additions and 11 deletions
+19 -4
View File
@@ -6,6 +6,7 @@ from typing import Iterable, List, Optional
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from app.database.models import WebApiToken
@@ -14,8 +15,12 @@ async def list_tokens(
db: AsyncSession,
*,
include_inactive: bool = False,
user_id: Optional[int] = None,
) -> List[WebApiToken]:
query = select(WebApiToken)
query = select(WebApiToken).options(selectinload(WebApiToken.user))
if user_id is not None:
query = query.where(WebApiToken.user_id == user_id)
if not include_inactive:
query = query.where(WebApiToken.is_active.is_(True))
@@ -27,12 +32,20 @@ async def list_tokens(
async def get_token_by_id(db: AsyncSession, token_id: int) -> Optional[WebApiToken]:
return await db.get(WebApiToken, token_id)
query = (
select(WebApiToken)
.options(selectinload(WebApiToken.user))
.where(WebApiToken.id == token_id)
)
result = await db.execute(query)
return result.scalar_one_or_none()
async def get_token_by_hash(db: AsyncSession, token_hash: str) -> Optional[WebApiToken]:
query = select(WebApiToken).where(
WebApiToken.token_hash == token_hash
query = (
select(WebApiToken)
.options(selectinload(WebApiToken.user))
.where(WebApiToken.token_hash == token_hash)
)
result = await db.execute(query)
return result.scalar_one_or_none()
@@ -47,8 +60,10 @@ async def create_token(
description: Optional[str] = None,
expires_at: Optional[datetime] = None,
created_by: Optional[str] = None,
user_id: Optional[int] = None,
) -> WebApiToken:
token = WebApiToken(
user_id=user_id,
name=name,
token_hash=token_hash,
token_prefix=token_prefix,
+3
View File
@@ -387,6 +387,7 @@ class User(Base):
referral_earnings = relationship("ReferralEarning", foreign_keys="ReferralEarning.user_id", back_populates="user")
discount_offers = relationship("DiscountOffer", back_populates="user")
promo_offer_logs = relationship("PromoOfferLog", back_populates="user")
web_api_tokens = relationship("WebApiToken", back_populates="user")
lifetime_used_traffic_bytes = Column(BigInteger, default=0)
auto_promo_group_assigned = Column(Boolean, nullable=False, default=False)
auto_promo_group_threshold_kopeks = Column(BigInteger, nullable=False, default=0)
@@ -1238,6 +1239,7 @@ class WebApiToken(Base):
__tablename__ = "web_api_tokens"
id = Column(Integer, primary_key=True, index=True)
user_id = Column(Integer, ForeignKey("users.id", ondelete="SET NULL"), nullable=True, index=True)
name = Column(String(255), nullable=False)
token_hash = Column(String(128), nullable=False, unique=True, index=True)
token_prefix = Column(String(32), nullable=False, index=True)
@@ -1249,6 +1251,7 @@ class WebApiToken(Base):
last_used_ip = Column(String(64), nullable=True)
is_active = Column(Boolean, default=True, nullable=False)
created_by = Column(String(255), nullable=True)
user = relationship("User", back_populates="web_api_tokens")
def __repr__(self) -> str:
status = "active" if self.is_active else "revoked"
+77 -3
View File
@@ -2434,6 +2434,7 @@ async def create_web_api_tokens_table() -> bool:
create_sql = """
CREATE TABLE web_api_tokens (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NULL,
name VARCHAR(255) NOT NULL,
token_hash VARCHAR(128) NOT NULL UNIQUE,
token_prefix VARCHAR(32) NOT NULL,
@@ -2444,16 +2445,19 @@ async def create_web_api_tokens_table() -> bool:
last_used_at DATETIME NULL,
last_used_ip VARCHAR(64) NULL,
is_active BOOLEAN NOT NULL DEFAULT 1,
created_by VARCHAR(255) NULL
created_by VARCHAR(255) NULL,
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL
);
CREATE INDEX idx_web_api_tokens_active ON web_api_tokens(is_active);
CREATE INDEX idx_web_api_tokens_prefix ON web_api_tokens(token_prefix);
CREATE INDEX idx_web_api_tokens_last_used ON web_api_tokens(last_used_at);
CREATE INDEX ix_web_api_tokens_user_id ON web_api_tokens(user_id);
"""
elif db_type == "postgresql":
create_sql = """
CREATE TABLE web_api_tokens (
id SERIAL PRIMARY KEY,
user_id INTEGER NULL,
name VARCHAR(255) NOT NULL,
token_hash VARCHAR(128) NOT NULL UNIQUE,
token_prefix VARCHAR(32) NOT NULL,
@@ -2464,16 +2468,19 @@ async def create_web_api_tokens_table() -> bool:
last_used_at TIMESTAMP NULL,
last_used_ip VARCHAR(64) NULL,
is_active BOOLEAN NOT NULL DEFAULT TRUE,
created_by VARCHAR(255) NULL
created_by VARCHAR(255) NULL,
CONSTRAINT fk_web_api_tokens_user_id FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL
);
CREATE INDEX idx_web_api_tokens_active ON web_api_tokens(is_active);
CREATE INDEX idx_web_api_tokens_prefix ON web_api_tokens(token_prefix);
CREATE INDEX idx_web_api_tokens_last_used ON web_api_tokens(last_used_at);
CREATE INDEX ix_web_api_tokens_user_id ON web_api_tokens(user_id);
"""
else:
create_sql = """
CREATE TABLE web_api_tokens (
id INT AUTO_INCREMENT PRIMARY KEY,
user_id INT NULL,
name VARCHAR(255) NOT NULL,
token_hash VARCHAR(128) NOT NULL UNIQUE,
token_prefix VARCHAR(32) NOT NULL,
@@ -2484,7 +2491,8 @@ async def create_web_api_tokens_table() -> bool:
last_used_at TIMESTAMP NULL,
last_used_ip VARCHAR(64) NULL,
is_active BOOLEAN NOT NULL DEFAULT TRUE,
created_by VARCHAR(255) NULL
created_by VARCHAR(255) NULL,
CONSTRAINT fk_web_api_tokens_user_id FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL
) ENGINE=InnoDB;
CREATE INDEX idx_web_api_tokens_active ON web_api_tokens(is_active);
CREATE INDEX idx_web_api_tokens_prefix ON web_api_tokens(token_prefix);
@@ -2500,6 +2508,71 @@ async def create_web_api_tokens_table() -> bool:
return False
async def ensure_web_api_tokens_user_column() -> bool:
column_exists = await check_column_exists("web_api_tokens", "user_id")
index_name = "ix_web_api_tokens_user_id"
constraint_name = "fk_web_api_tokens_user_id"
try:
async with engine.begin() as conn:
db_type = await get_database_type()
if not column_exists:
if db_type == "sqlite":
await conn.execute(text("ALTER TABLE web_api_tokens ADD COLUMN user_id INTEGER"))
elif db_type == "postgresql":
await conn.execute(text("ALTER TABLE web_api_tokens ADD COLUMN user_id INTEGER"))
elif db_type == "mysql":
await conn.execute(text("ALTER TABLE web_api_tokens ADD COLUMN user_id INT NULL"))
else:
logger.error(f"Неподдерживаемый тип БД для web_api_tokens.user_id: {db_type}")
return False
logger.info("Добавлена колонка web_api_tokens.user_id")
if db_type in {"postgresql", "mysql"}:
constraint_exists = await check_constraint_exists("web_api_tokens", constraint_name)
if not constraint_exists:
try:
await conn.execute(
text(
"ALTER TABLE web_api_tokens "
"ADD CONSTRAINT fk_web_api_tokens_user_id "
"FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE SET NULL"
)
)
logger.info("Добавлено ограничение fk_web_api_tokens_user_id")
except Exception as constraint_error:
logger.warning(
"Не удалось создать ограничение fk_web_api_tokens_user_id: %s",
constraint_error,
)
index_exists = await check_index_exists("web_api_tokens", index_name)
if not index_exists:
try:
async with engine.begin() as conn:
db_type = await get_database_type()
if db_type == "sqlite":
await conn.execute(
text("CREATE INDEX IF NOT EXISTS ix_web_api_tokens_user_id ON web_api_tokens(user_id)")
)
elif db_type == "postgresql":
await conn.execute(
text("CREATE INDEX IF NOT EXISTS ix_web_api_tokens_user_id ON web_api_tokens(user_id)")
)
elif db_type == "mysql":
await conn.execute(text("CREATE INDEX ix_web_api_tokens_user_id ON web_api_tokens(user_id)"))
logger.info("Создан индекс ix_web_api_tokens_user_id")
except Exception as index_error:
logger.warning(f"Не удалось создать индекс ix_web_api_tokens_user_id: {index_error}")
return True
except Exception as error:
logger.error(f"Ошибка обновления web_api_tokens.user_id: {error}")
return False
async def create_privacy_policies_table() -> bool:
table_exists = await check_table_exists("privacy_policies")
if table_exists:
@@ -2799,6 +2872,7 @@ async def run_universal_migration():
web_api_tokens_ready = await create_web_api_tokens_table()
if web_api_tokens_ready:
logger.info("✅ Таблица web_api_tokens готова")
await ensure_web_api_tokens_user_column()
else:
logger.warning("⚠️ Проблемы с таблицей web_api_tokens")
+159 -1
View File
@@ -12,10 +12,12 @@ from app.keyboards.inline import (
get_main_menu_keyboard,
get_language_selection_keyboard,
get_info_menu_keyboard,
get_api_access_keyboard,
)
from app.localization.texts import get_texts, get_rules
from app.database.models import User
from app.database.models import User, WebApiToken
from app.database.crud.user_message import get_random_active_message
from app.database.crud import web_api_token as web_api_token_crud
from app.services.subscription_checkout_service import (
has_subscription_checkout_draft,
should_offer_checkout_resume,
@@ -29,6 +31,8 @@ from app.utils.promo_offer import (
from app.services.privacy_policy_service import PrivacyPolicyService
from app.services.public_offer_service import PublicOfferService
from app.services.faq_service import FaqService
from app.services.web_api_token_service import web_api_token_service
from app.utils.formatters import format_datetime
logger = logging.getLogger(__name__)
@@ -132,6 +136,150 @@ async def show_info_menu(
await callback.answer()
def _build_api_access_caption(
texts,
token: WebApiToken | None,
*,
has_active_token: bool,
) -> str:
header = texts.t("API_ACCESS_HEADER", "🔑 <b>API доступ</b>")
parts: list[str] = [header]
if not token:
parts.append(
texts.t(
"API_ACCESS_NO_TOKEN",
(
"У вас еще нет активного API ключа.\n"
"Нажмите кнопку ниже, чтобы выпустить первый ключ."
),
)
)
else:
created_at = format_datetime(token.created_at) if token.created_at else ""
last_used = (
format_datetime(token.last_used_at)
if token.last_used_at
else texts.t("API_ACCESS_LAST_USED_NEVER", "еще не использовался")
)
last_ip_line = ""
if token.last_used_ip:
last_ip_line = "\n" + texts.t(
"API_ACCESS_LAST_IP_LINE",
"Последний IP: {ip}",
).format(ip=token.last_used_ip)
template_key = "API_ACCESS_ACTIVE_TOKEN" if has_active_token else "API_ACCESS_INACTIVE_TOKEN"
parts.append(
texts.t(
template_key,
(
"Ваш текущий API ключ выпущен {created_at}.\n"
"Первые символы: <code>{prefix}</code>.\n"
"Последнее использование: {last_used}.{last_ip_line}\n"
"Чтобы выпустить новый ключ, нажмите кнопку ниже — предыдущий будет отключен."
)
if has_active_token
else (
"Последний ключ выпущен {created_at}.\n"
"Первые символы: <code>{prefix}</code>.\n"
"Последнее использование: {last_used}.{last_ip_line}\n"
"Чтобы получить доступ, создайте новый ключ."
),
).format(
created_at=created_at,
prefix=token.token_prefix,
last_used=last_used,
last_ip_line=last_ip_line,
)
)
parts.append(
texts.t(
"API_ACCESS_SECURITY_HINT",
"Берегите ключ и не передавайте его другим людям.",
)
)
return "\n\n".join(part for part in parts if part)
async def show_api_access(
callback: types.CallbackQuery,
db_user: User,
db: AsyncSession,
*,
skip_callback_answer: bool = False,
) -> None:
texts = get_texts(db_user.language)
tokens = await web_api_token_crud.list_tokens(
db,
include_inactive=True,
user_id=db_user.id,
)
active_token = next((token for token in tokens if token.is_active), None)
latest_token = tokens[0] if tokens else None
token_to_show = active_token or latest_token
caption = _build_api_access_caption(
texts,
token_to_show,
has_active_token=bool(active_token),
)
await edit_or_answer_photo(
callback=callback,
caption=caption,
keyboard=get_api_access_keyboard(db_user.language),
parse_mode="HTML",
)
if not skip_callback_answer:
await callback.answer()
async def generate_api_access_token(
callback: types.CallbackQuery,
db_user: User,
db: AsyncSession,
) -> None:
texts = get_texts(db_user.language)
token_value, token = await web_api_token_service.issue_user_token(db, db_user)
await db.commit()
logger.info(
"🔑 Пользователь %s выпустил новый API ключ (token_id=%s)",
db_user.telegram_id,
token.id,
)
message_text = texts.t(
"API_ACCESS_NEW_TOKEN",
(
"Ваш новый API ключ:\n"
"<code>{token}</code>\n\n"
"Сохраните его — повторно показать не получится."
),
).format(token=token_value)
await callback.message.answer(message_text, parse_mode="HTML")
await show_api_access(
callback,
db_user,
db,
skip_callback_answer=True,
)
await callback.answer(
texts.t("API_ACCESS_GENERATED_ALERT", "Новый API ключ создан"),
show_alert=True,
)
async def show_faq_pages(
callback: types.CallbackQuery,
db_user: User,
@@ -782,6 +930,16 @@ def register_handlers(dp: Dispatcher):
F.data == "menu_info",
)
dp.callback_query.register(
show_api_access,
F.data == "menu_api_access",
)
dp.callback_query.register(
generate_api_access_token,
F.data == "api_access_generate",
)
dp.callback_query.register(
show_faq_pages,
F.data == "menu_faq",
+23
View File
@@ -314,6 +314,13 @@ def get_info_menu_keyboard(
)
])
buttons.append([
InlineKeyboardButton(
text=texts.t("MENU_API_ACCESS", "🔑 API доступ"),
callback_data="menu_api_access",
)
])
buttons.append([
InlineKeyboardButton(text=texts.MENU_RULES, callback_data="menu_rules")
])
@@ -347,6 +354,22 @@ def get_info_menu_keyboard(
return InlineKeyboardMarkup(inline_keyboard=buttons)
def get_api_access_keyboard(language: str = DEFAULT_LANGUAGE) -> InlineKeyboardMarkup:
texts = get_texts(language)
buttons = [
[
InlineKeyboardButton(
text=texts.t("API_ACCESS_GENERATE_BUTTON", "🔄 Выпустить ключ"),
callback_data="api_access_generate",
)
],
[InlineKeyboardButton(text=texts.BACK, callback_data="menu_info")],
]
return InlineKeyboardMarkup(inline_keyboard=buttons)
def get_happ_download_button_row(texts) -> Optional[List[InlineKeyboardButton]]:
if not settings.is_happ_download_button_enabled():
return None
+11
View File
@@ -305,6 +305,7 @@
"MENU_INFO": "️ Info",
"MENU_INFO_HEADER": "️ <b>Info</b>",
"MENU_INFO_PROMPT": "Choose a section:",
"MENU_API_ACCESS": "🔑 API key",
"MENU_PROMOCODE": "🎫 Promo code",
"MENU_REFERRALS": "🤝 Referral program",
"MENU_RULES": "📋 Service rules",
@@ -759,3 +760,13 @@
"ADMIN_PUBLIC_OFFER_RETURN_TO_EDIT": "⬅️ Back to editing",
"ADMIN_PUBLIC_OFFER_EDIT_CANCELLED": "Offer editing cancelled."
}
"API_ACCESS_HEADER": "🔑 <b>API access</b>",
"API_ACCESS_NO_TOKEN": "You don't have an active API key yet.\nTap the button below to issue your first key.\n\nEach new key disables the previous one.",
"API_ACCESS_LAST_USED_NEVER": "not used yet",
"API_ACCESS_LAST_IP_LINE": "Last IP: {ip}",
"API_ACCESS_ACTIVE_TOKEN": "Your current API key was issued on {created_at}.\nFirst characters: <code>{prefix}</code>.\nLast used: {last_used}.{last_ip_line}\nTap the button below to issue a new key — the previous one will be revoked.",
"API_ACCESS_INACTIVE_TOKEN": "The last key was issued on {created_at}.\nFirst characters: <code>{prefix}</code>.\nLast used: {last_used}.{last_ip_line}\nIssue a new key to regain access — old keys no longer work.",
"API_ACCESS_SECURITY_HINT": "Store the key securely and never share it with anyone.",
"API_ACCESS_GENERATE_BUTTON": "🔄 Issue key",
"API_ACCESS_NEW_TOKEN": "Here is your new API key:\n<code>{token}</code>\n\nSave it now — it cannot be shown again. Previous keys have been revoked.",
"API_ACCESS_GENERATED_ALERT": "New API key issued",
+11
View File
@@ -236,6 +236,17 @@
"MENU_INFO_PROMPT": "Выберите раздел:",
"MENU_PRIVACY_POLICY": "🛡️ Политика конф.",
"MENU_PUBLIC_OFFER": "📄 Оферта",
"MENU_API_ACCESS": "🔑 API доступ",
"API_ACCESS_HEADER": "🔑 <b>API доступ</b>",
"API_ACCESS_NO_TOKEN": "У вас ещё нет активного API ключа.\nНажмите кнопку ниже, чтобы выпустить первый ключ.\n\nКаждый новый ключ отключает предыдущие.",
"API_ACCESS_LAST_USED_NEVER": "ещё не использовался",
"API_ACCESS_LAST_IP_LINE": "Последний IP: {ip}",
"API_ACCESS_ACTIVE_TOKEN": "Ваш текущий API ключ выпущен {created_at}.\nПервые символы: <code>{prefix}</code>.\nПоследнее использование: {last_used}.{last_ip_line}\nЧтобы выпустить новый ключ, нажмите кнопку ниже — предыдущий будет отключен.",
"API_ACCESS_INACTIVE_TOKEN": "Последний ключ выпущен {created_at}.\nПервые символы: <code>{prefix}</code>.\nПоследнее использование: {last_used}.{last_ip_line}\nЧтобы получить доступ, создайте новый ключ — старые больше не работают.",
"API_ACCESS_SECURITY_HINT": "Храните ключ в безопасном месте и не передавайте его другим людям.",
"API_ACCESS_GENERATE_BUTTON": "🔄 Выпустить ключ",
"API_ACCESS_NEW_TOKEN": "Вот ваш новый API ключ:\n<code>{token}</code>\n\nСохраните его сейчас — повторно показать не получится. Предыдущие ключи отключены.",
"API_ACCESS_GENERATED_ALERT": "Новый API ключ создан",
"PRIVACY_POLICY_HEADER": "🛡️ <b>Политика конфиденциальности</b>",
"PRIVACY_POLICY_NOT_AVAILABLE": "Политика конфиденциальности временно недоступна.",
"PRIVACY_POLICY_EMPTY_ALERT": "Политика конфиденциальности ещё не заполнена.",
+42 -1
View File
@@ -8,7 +8,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from app.config import settings
from app.database.crud import web_api_token as crud
from app.database.models import WebApiToken
from app.database.models import User, WebApiToken
from app.database.universal_migration import ensure_default_web_api_token
from app.utils.security import generate_api_token, hash_api_token
@@ -66,10 +66,14 @@ class WebApiTokenService:
expires_at: Optional[datetime] = None,
created_by: Optional[str] = None,
token_value: Optional[str] = None,
user: Optional[User] = None,
user_id: Optional[int] = None,
) -> Tuple[str, WebApiToken]:
plain_token = token_value or generate_api_token()
token_hash = self.hash_token(plain_token)
resolved_user_id = user.id if user else user_id
token = await crud.create_token(
db,
name=name,
@@ -78,6 +82,43 @@ class WebApiTokenService:
description=description,
expires_at=expires_at,
created_by=created_by,
user_id=resolved_user_id,
)
return plain_token, token
async def issue_user_token(
self,
db: AsyncSession,
user: User,
) -> Tuple[str, WebApiToken]:
existing_tokens = await crud.list_tokens(
db,
include_inactive=True,
user_id=user.id,
)
now = datetime.utcnow()
has_updates = False
for token in existing_tokens:
if token.is_active:
token.is_active = False
token.updated_at = now
has_updates = True
if has_updates:
await db.flush()
token_name = f"User {user.telegram_id}"
description = "Generated via Telegram bot"
created_by = f"telegram:{user.telegram_id}"
plain_token, token = await self.create_token(
db,
name=token_name,
description=description,
created_by=created_by,
user=user,
)
return plain_token, token
+19 -2
View File
@@ -1,8 +1,11 @@
from __future__ import annotations
from fastapi import APIRouter, Depends, HTTPException, Response, Security, status
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query, Response, Security, status
from sqlalchemy.ext.asyncio import AsyncSession
from app.database.crud.user import get_user_by_id
from app.database.crud.web_api_token import (
delete_token,
get_token_by_id,
@@ -30,6 +33,8 @@ def _serialize(token: WebApiToken) -> TokenResponse:
last_used_at=token.last_used_at,
last_used_ip=token.last_used_ip,
created_by=token.created_by,
user_id=token.user_id,
user_telegram_id=getattr(token.user, "telegram_id", None),
)
@@ -37,8 +42,13 @@ def _serialize(token: WebApiToken) -> TokenResponse:
async def get_tokens(
_: WebApiToken = Security(require_api_token),
db: AsyncSession = Depends(get_db_session),
user_id: Optional[int] = Query(default=None, description="Фильтр по ID пользователя"),
) -> list[TokenResponse]:
tokens = await list_tokens(db, include_inactive=True)
tokens = await list_tokens(
db,
include_inactive=True,
user_id=user_id,
)
return [_serialize(token) for token in tokens]
@@ -48,12 +58,19 @@ async def create_token(
actor: WebApiToken = Security(require_api_token),
db: AsyncSession = Depends(get_db_session),
) -> TokenCreateResponse:
target_user = None
if payload.user_id is not None:
target_user = await get_user_by_id(db, payload.user_id)
if not target_user:
raise HTTPException(status.HTTP_404_NOT_FOUND, "User not found")
token_value, token = await web_api_token_service.create_token(
db,
name=payload.name.strip(),
description=payload.description,
expires_at=payload.expires_at,
created_by=actor.name,
user=target_user,
)
await db.commit()
+12
View File
@@ -18,12 +18,24 @@ class TokenResponse(BaseModel):
last_used_at: Optional[datetime] = None
last_used_ip: Optional[str] = None
created_by: Optional[str] = None
user_id: Optional[int] = Field(
default=None,
description="ID пользователя, которому принадлежит токен",
)
user_telegram_id: Optional[int] = Field(
default=None,
description="Telegram ID владельца токена",
)
class TokenCreateRequest(BaseModel):
name: str
description: Optional[str] = None
expires_at: Optional[datetime] = None
user_id: Optional[int] = Field(
default=None,
description="ID пользователя, для которого выпускается токен",
)
class TokenCreateResponse(TokenResponse):
@@ -0,0 +1,72 @@
"""add user link to web api tokens"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.engine import reflection
revision: str = "d2a9443a9f47"
down_revision: Union[str, None] = "8fd1e338eb45"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
TABLE_NAME = "web_api_tokens"
COLUMN_NAME = "user_id"
FK_NAME = "fk_web_api_tokens_user_id"
INDEX_NAME = "ix_web_api_tokens_user_id"
def _column_exists(inspector: reflection.Inspector) -> bool:
return COLUMN_NAME in [column["name"] for column in inspector.get_columns(TABLE_NAME)]
def _fk_exists(inspector: reflection.Inspector) -> bool:
return any(fk.get("name") == FK_NAME for fk in inspector.get_foreign_keys(TABLE_NAME))
def _index_exists(inspector: reflection.Inspector) -> bool:
return any(index.get("name") == INDEX_NAME for index in inspector.get_indexes(TABLE_NAME))
def upgrade() -> None:
bind = op.get_bind()
inspector = reflection.Inspector.from_engine(bind)
if not _column_exists(inspector):
op.add_column(TABLE_NAME, sa.Column(COLUMN_NAME, sa.Integer(), nullable=True))
inspector = reflection.Inspector.from_engine(bind)
if bind.dialect.name in {"postgresql", "mysql"} and not _fk_exists(inspector):
try:
op.create_foreign_key(
FK_NAME,
TABLE_NAME,
"users",
[COLUMN_NAME],
["id"],
ondelete="SET NULL",
)
except Exception:
# Constraint creation can fail if duplicates exist; skip to keep migration resilient.
pass
inspector = reflection.Inspector.from_engine(bind)
if not _index_exists(inspector):
op.create_index(INDEX_NAME, TABLE_NAME, [COLUMN_NAME])
def downgrade() -> None:
bind = op.get_bind()
inspector = reflection.Inspector.from_engine(bind)
if _index_exists(inspector):
op.drop_index(INDEX_NAME, table_name=TABLE_NAME)
inspector = reflection.Inspector.from_engine(bind)
if _fk_exists(inspector):
op.drop_constraint(FK_NAME, TABLE_NAME, type_="foreignkey")
inspector = reflection.Inspector.from_engine(bind)
if _column_exists(inspector):
op.drop_column(TABLE_NAME, COLUMN_NAME)