Method get_next_available_credential

src/storage/psql_manager.py:245–319

随机获取一个可用凭证（负载均衡）

(
        self, mode: str = "geminicli", model_name: Optional[str] = None
    )

    # ============ 凭证查询方法 ============

    async def get_next_available_credential(
        self, mode: str = "geminicli", model_name: Optional[str] = None
    ) -> Optional[Tuple[str, Dict[str, Any]]]:
        """随机获取一个可用凭证（负载均衡）"""
        self._ensure_initialized()

        try:
            table_name = self._get_table_name(mode)
            current_time = time.time()

            async with self._pool.acquire() as conn:
                if mode == "geminicli":
                    rows = await conn.fetch(f"""
                        SELECT filename, credential_data, model_cooldowns, preview
                        FROM {table_name}
                        WHERE disabled = 0
                        ORDER BY RANDOM()
                    """)

                    if not model_name:
                        if rows:
                            return rows[0]["filename"], json.loads(rows[0]["credential_data"])
                        return None

                    is_preview_model = "preview" in model_name.lower()
                    non_preview_creds = []
                    preview_creds = []

                    for row in rows:
                        model_cooldowns = json.loads(row["model_cooldowns"] or "{}")
                        cd = model_cooldowns.get(model_name)
                        if cd is None or current_time >= cd:
                            if row["preview"]:
                                preview_creds.append((row["filename"], row["credential_data"]))
                            else:
                                non_preview_creds.append((row["filename"], row["credential_data"]))

                    if is_preview_model:
                        if preview_creds:
                            return preview_creds[0][0], json.loads(preview_creds[0][1])
                    else:
                        if non_preview_creds:
                            return non_preview_creds[0][0], json.loads(non_preview_creds[0][1])
                        elif preview_creds:
                            return preview_creds[0][0], json.loads(preview_creds[0][1])

                    return None
                else:
                    rows = await conn.fetch(f"""
                        SELECT filename, credential_data, model_cooldowns, enable_credit
                        FROM {table_name}
                        WHERE disabled = 0
                        ORDER BY RANDOM()
                    """)

                    if not model_name:
                        if rows:
                            credential_data = json.loads(rows[0]["credential_data"])