aboutsummaryrefslogtreecommitdiffstats
path: root/arthur/exts/kubernetes/pods.py
blob: 0a82fb999dc72262e39522cc9b10b3de34da7ec5 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
"""The Pods cog helps with managing Kubernetes pods."""

import zoneinfo
from datetime import datetime

import humanize
from discord.ext import commands
from kubernetes_asyncio.client.rest import ApiException
from loguru import logger
from tabulate import tabulate

from arthur.apis.kubernetes import pods
from arthur.bot import KingArthur
from arthur.config import CONFIG
from arthur.utils import generate_error_message

MAX_MESSAGE_LENGTH = 2000


def tabulate_pod_data(data: list[list[str]]) -> str:
    """Tabulate the pod data to be sent to Discord."""
    table = tabulate(
        data,
        headers=["Status", "Pod", "Phase", "IP", "Node", "Age", "Restarts"],
        tablefmt="psql",
        colalign=("center", "left", "left", "center", "center", "left", "center"),
    )

    return f"```\n{table}```"


class Pods(commands.Cog):
    """Commands for working with Kubernetes Pods."""

    def __init__(self, bot: KingArthur) -> None:
        self.bot = bot

    @commands.group(name="pods", aliases=["pod"], invoke_without_command=True)
    async def pods_cmd(self, ctx: commands.Context) -> None:
        """Commands for working with Kubernetes Pods."""
        await ctx.send_help(ctx.command)

    @pods_cmd.command(name="list", aliases=["ls"])
    async def pods_list(self, ctx: commands.Context, namespace: str = "default") -> None:
        """List pods in the selected namespace (defaults to default)."""
        pod_list = await pods.list_pods(namespace)

        if len(pod_list.items) == 0:
            return await ctx.send(
                generate_error_message(description="No pods found, check the namespace exists.")
            )

        tables = [[]]

        for pod in pod_list.items:
            match pod.status.phase:
                case "Running":
                    emote = "\N{LARGE GREEN CIRCLE}"
                case "Pending":
                    emote = "\N{LARGE YELLOW CIRCLE}"
                case "Succeeded":
                    emote = "\N{WHITE HEAVY CHECK MARK}"
                case "Failed":
                    emote = "\N{CROSS MARK}"
                case "Unknown":
                    emote = "\N{WHITE QUESTION MARK ORNAMENT}"
                case _:
                    emote = "\N{BLACK QUESTION MARK ORNAMENT}"

            time_human = humanize.naturaldelta(
                datetime.now(tz=zoneinfo.ZoneInfo("UTC")) - pod.metadata.creation_timestamp
            )

            # we know that Linode formats names like "lke<cluster>-<pool>-<node>"
            node_name = pod.spec.node_name.split("-")[2]

            table_data = [
                emote,
                pod.metadata.name,
                pod.status.phase,
                pod.status.pod_ip,
                node_name,
                time_human,
                pod.status.container_statuses[0].restart_count,
            ]

            if len(tabulate_pod_data(tables[-1] + [table_data])) > MAX_MESSAGE_LENGTH:
                tables.append([])
                tables[-1].append(table_data)
            else:
                tables[-1].append(table_data)

        await ctx.send(f"**Pods in namespace `{namespace}`**")

        for table in tables:
            await ctx.send(tabulate_pod_data(table))

        return None

    @pods_cmd.command(name="logs", aliases=["log", "tail"])
    @commands.check(lambda ctx: ctx.channel.id == CONFIG.devops_channel_id)
    async def pods_logs(
        self, ctx: commands.Context, pod_name: str, namespace: str = "default", lines: int = 15
    ) -> None:
        """
        Tail the logs of a pod in the selected namespace (defaults to default).

        We also support the syntax of `deploy/<deployment-name>` to get the logs of the first pod associated with the deployment.
        """
        if pod_name.startswith("deploy/"):
            pod_names = await pods.get_pod_names_from_deployment(
                namespace, pod_name.removeprefix("deploy/")
            )
            logger.debug(f"Resolved deployment pod name to {pod_names}")
        else:
            pod_names = [pod_name]

        if pod_names is None:
            return await ctx.send(
                generate_error_message(description="No pods found for the provided deployment.")
            )

        for pod in pod_names:
            try:
                logs = await pods.tail_pod(namespace, pod, lines=lines)
            except ApiException as e:
                if e.status == 404:  # noqa: PLR2004, 404 is a known error
                    return await ctx.send(
                        generate_error_message(
                            description="Pod or namespace not found, check the name."
                        )
                    )
                return await ctx.send(generate_error_message(description=str(e)))

            if len(logs) == 0:
                return await ctx.send(
                    generate_error_message(description="No logs found for the pod.")
                )

            truncated = False

            if len(logs) > MAX_MESSAGE_LENGTH - 100:
                truncated = True
                while len(logs) > MAX_MESSAGE_LENGTH - 100:
                    logs = logs[: logs.rfind("\n")]

            message = f"**Logs for pod `{pod}` in namespace `{namespace}`**\n"

            if truncated:
                message += "`[Logs truncated]`\n"

            message += "```"
            message += logs
            message += "```"

            await ctx.send(message)

        return None


async def setup(bot: KingArthur) -> None:
    """Add the extension to the bot."""
    await bot.add_cog(Pods(bot))