1# The following code and documentation was inspired, and in some cases
2# copied and modified, from the work of Vinay Sajip and contributors
3# on cpython's logging package
4
5import abc
6import asyncio
7import datetime
8import enum
9import os
10import re
11import time
12from asyncio import AbstractEventLoop
13from typing import Callable, List, Optional
14
15import aiofiles
16from aiofiles.threadpool import AsyncTextIOWrapper
17
18from aiologger.handlers.base import Handler
19from aiologger.records import LogRecord
20from aiologger.utils import classproperty, get_running_loop, loop_compat
21
22
23@loop_compat
24class AsyncFileHandler(Handler):
25    terminator = "\n"
26
27    def __init__(
28        self, filename: str, mode: str = "a", encoding: str = None
29    ) -> None:
30        super().__init__()
31        filename = os.fspath(filename)
32        self.absolute_file_path = os.path.abspath(filename)
33        self.mode = mode
34        self.encoding = encoding
35        self.stream: AsyncTextIOWrapper = None
36        self._initialization_lock = None
37
38    @property
39    def initialized(self):
40        return self.stream is not None
41
42    async def _init_writer(self):
43        """
44        Open the current base file with the (original) mode and encoding.
45        """
46        if not self._initialization_lock:
47            self._initialization_lock = asyncio.Lock()
48
49        async with self._initialization_lock:
50            if not self.initialized:
51                self.stream = await aiofiles.open(
52                    file=self.absolute_file_path,
53                    mode=self.mode,
54                    encoding=self.encoding,
55                )
56
57    async def flush(self):
58        await self.stream.flush()
59
60    async def close(self):
61        if not self.initialized:
62            return
63        await self.stream.flush()
64        await self.stream.close()
65        self.stream = None
66        self._initialization_lock = None
67
68    async def emit(self, record: LogRecord):
69        if not self.initialized:
70            await self._init_writer()
71
72        try:
73            msg = self.formatter.format(record)
74
75            # Write order is not guaranteed. String concatenation required
76            await self.stream.write(msg + self.terminator)
77
78            await self.stream.flush()
79        except Exception as exc:
80            await self.handle_error(record, exc)
81
82
83Namer = Callable[[str], str]
84Rotator = Callable[[str, str], None]
85
86
87class BaseAsyncRotatingFileHandler(AsyncFileHandler, metaclass=abc.ABCMeta):
88    def __init__(
89        self,
90        filename: str,
91        mode: str = "a",
92        encoding: str = None,
93        namer: Namer = None,
94        rotator: Rotator = None,
95    ) -> None:
96        super().__init__(filename, mode, encoding)
97        self.mode = mode
98        self.encoding = encoding
99        self.namer = namer
100        self.rotator = rotator
101        self._rollover_lock: Optional[asyncio.Lock] = None
102
103    def should_rollover(self, record: LogRecord) -> bool:
104        raise NotImplementedError
105
106    async def do_rollover(self):
107        raise NotImplementedError
108
109    async def emit(self, record: LogRecord):  # type: ignore
110        """
111        Emit a record.
112
113        Output the record to the file, catering for rollover as described
114        in `do_rollover`.
115        """
116        try:
117            if self.should_rollover(record):
118                if not self._rollover_lock:
119                    self._rollover_lock = asyncio.Lock()
120
121                async with self._rollover_lock:
122                    if self.should_rollover(record):
123                        await self.do_rollover()
124            await super().emit(record)
125        except Exception as exc:
126            await self.handle_error(record, exc)
127
128    def rotation_filename(self, default_name: str) -> str:
129        """
130        Modify the filename of a log file when rotating.
131
132        This is provided so that a custom filename can be provided.
133
134        :param default_name: The default name for the log file.
135        """
136        if self.namer is None:
137            return default_name
138
139        return self.namer(default_name)
140
141    async def rotate(self, source: str, dest: str):
142        """
143        When rotating, rotate the current log.
144
145        The default implementation calls the 'rotator' attribute of the
146        handler, if it's callable, passing the source and dest arguments to
147        it. If the attribute isn't callable (the default is None), the source
148        is simply renamed to the destination.
149
150        :param source: The source filename. This is normally the base
151                       filename, e.g. 'test.log'
152        :param dest:   The destination filename. This is normally
153                       what the source is rotated to, e.g. 'test.log.1'.
154        """
155        if self.rotator is None:
156            # logging issue 18940: A file may not have been created if delay is True.
157            loop = get_running_loop()
158            if await loop.run_in_executor(None, lambda: os.path.exists(source)):
159                await loop.run_in_executor(  # type: ignore
160                    None, lambda: os.rename(source, dest)
161                )
162        else:
163            self.rotator(source, dest)
164
165
166class RolloverInterval(str, enum.Enum):
167    SECONDS = "S"
168    MINUTES = "M"
169    HOURS = "H"
170    DAYS = "D"
171    MONDAYS = "W0"
172    TUESDAYS = "W1"
173    WEDNESDAYS = "W2"
174    THURSDAYS = "W3"
175    FRIDAYS = "W4"
176    SATURDAYS = "W5"
177    SUNDAYS = "W6"
178    MIDNIGHT = "MIDNIGHT"
179
180    @classproperty
181    def WEEK_DAYS(cls):
182        return (
183            cls.MONDAYS,
184            cls.TUESDAYS,
185            cls.WEDNESDAYS,
186            cls.THURSDAYS,
187            cls.FRIDAYS,
188            cls.SATURDAYS,
189            cls.SUNDAYS,
190        )
191
192
193ONE_MINUTE_IN_SECONDS = 60
194ONE_HOUR_IN_SECONDS = 60 * 60
195ONE_DAY_IN_SECONDS = ONE_HOUR_IN_SECONDS * 24
196ONE_WEEK_IN_SECONDS = 7 * ONE_DAY_IN_SECONDS
197
198
199class AsyncTimedRotatingFileHandler(BaseAsyncRotatingFileHandler):
200    """
201    Handler for logging to a file, rotating the log file at certain timed
202    intervals.
203
204    If `backup_count` is > 0, when rollover is done, no more than `backup_count`
205    files are kept - the oldest ones are deleted.
206
207    | when       | at_time behavior                                       |
208    |------------|--------------------------------------------------------|
209    | SECONDS    | at_time will be ignored                                |
210    | MINUTES    | -- // --                                               |
211    | HOURS      | -- // --                                               |
212    | DAYS       | at_time will be IGNORED. See also MIDNIGHT             |
213    | MONDAYS    | rotation happens every WEEK on MONDAY at ${at_time}    |
214    | TUESDAYS   | rotation happens every WEEK on TUESDAY at ${at_time}   |
215    | WEDNESDAYS | rotation happens every WEEK on WEDNESDAY at ${at_time} |
216    | THURSDAYS  | rotation happens every WEEK on THURSDAY at ${at_time}  |
217    | FRIDAYS    | rotation happens every WEEK on FRIDAY at ${at_time}    |
218    | SATURDAYS  | rotation happens every WEEK on SATURDAY at ${at_time}  |
219    | SUNDAYS    | rotation happens every WEEK on SUNDAY at ${at_time}    |
220    | MIDNIGHT   | rotation happens every DAY at ${at_time}               |
221    """
222
223    def __init__(
224        self,
225        filename: str,
226        when: RolloverInterval = RolloverInterval.HOURS,
227        interval: int = 1,
228        backup_count: int = 0,
229        encoding: str = None,
230        utc: bool = False,
231        at_time: datetime.time = None,
232    ) -> None:
233        super().__init__(filename=filename, mode="a", encoding=encoding)
234        self.when = when.upper()
235        self.backup_count = backup_count
236        self.utc = utc
237        self.at_time = at_time
238        # Calculate the real rollover interval, which is just the number of
239        # seconds between rollovers.  Also set the filename suffix used when
240        # a rollover occurs.  Current 'when' events supported:
241        # S - Seconds
242        # M - Minutes
243        # H - Hours
244        # D - Days
245        # midnight - roll over at midnight
246        # W{0-6} - roll over on a certain day; 0 - Monday
247        #
248        # Case of the 'when' specifier is not important; lower or upper case
249        # will work.
250        if self.when == RolloverInterval.SECONDS:
251            self.interval = 1  # one second
252            self.suffix = "%Y-%m-%d_%H-%M-%S"
253            ext_match = r"^\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}(\.\w+)?$"
254        elif self.when == RolloverInterval.MINUTES:
255            self.interval = ONE_MINUTE_IN_SECONDS  # one minute
256            self.suffix = "%Y-%m-%d_%H-%M"
257            ext_match = r"^\d{4}-\d{2}-\d{2}_\d{2}-\d{2}(\.\w+)?$"
258        elif self.when == RolloverInterval.HOURS:
259            self.interval = ONE_HOUR_IN_SECONDS  # one hour
260            self.suffix = "%Y-%m-%d_%H"
261            ext_match = r"^\d{4}-\d{2}-\d{2}_\d{2}(\.\w+)?$"
262        elif (
263            self.when == RolloverInterval.DAYS
264            or self.when == RolloverInterval.MIDNIGHT
265        ):
266            self.interval = ONE_DAY_IN_SECONDS  # one day
267            self.suffix = "%Y-%m-%d"
268            ext_match = r"^\d{4}-\d{2}-\d{2}(\.\w+)?$"
269        elif self.when.startswith("W"):
270            if self.when not in RolloverInterval.WEEK_DAYS:
271                raise ValueError(
272                    f"Invalid day specified for weekly rollover: {self.when}"
273                )
274            self.interval = ONE_DAY_IN_SECONDS * 7  # one week
275            self.day_of_week = int(self.when[1])
276            self.suffix = "%Y-%m-%d"
277            ext_match = r"^\d{4}-\d{2}-\d{2}(\.\w+)?$"
278        else:
279            raise ValueError(f"Invalid RolloverInterval specified: {self.when}")
280
281        self.ext_match = re.compile(ext_match, re.ASCII)
282        self.interval = self.interval * interval  # multiply by units requested
283        # The following line added because the filename passed in could be a
284        # path object (see Issue #27493), but self.baseFilename will be a string
285        filename = self.absolute_file_path
286        if os.path.exists(filename):  # todo: IO. Remove or postpone
287            t = int(os.stat(filename).st_mtime)
288        else:
289            t = int(time.time())
290        self.rollover_at = self.compute_rollover(t)
291
292    def compute_rollover(self, current_time: int) -> int:
293        """
294        Work out the rollover time based on the specified time.
295
296        If we are rolling over at midnight or weekly, then the interval is
297        already known. need to figure out is WHEN the next interval is.
298        In other words, if you are rolling over at midnight, then your base
299        interval is 1 day, but you want to start that one day clock at midnight,
300         not now. So, we have to fudge the `rollover_at` value in order to trigger
301         the first rollover at the right time.  After that, the regular interval
302         will take care of the rest.  Note that this code doesn't care about
303         leap seconds. :)
304        """
305        result = current_time + self.interval
306
307        if (
308            self.when == RolloverInterval.MIDNIGHT
309            or self.when in RolloverInterval.WEEK_DAYS
310        ):
311            if self.utc:
312                t = time.gmtime(current_time)
313            else:
314                t = time.localtime(current_time)
315            current_hour = t[3]
316            current_minute = t[4]
317            current_second = t[5]
318            current_day = t[6]
319            # r is the number of seconds left between now and the next rotation
320            if self.at_time is None:
321                rotate_ts = ONE_DAY_IN_SECONDS
322            else:
323                rotate_ts = (
324                    self.at_time.hour * 60 + self.at_time.minute
325                ) * 60 + self.at_time.second
326
327            r = rotate_ts - (
328                (current_hour * 60 + current_minute) * 60 + current_second
329            )
330            if r < 0:
331                # Rotate time is before the current time (for example when
332                # self.rotateAt is 13:45 and it now 14:15), rotation is
333                # tomorrow.
334                r += ONE_DAY_IN_SECONDS
335                current_day = (current_day + 1) % 7
336            result = current_time + r
337            # If we are rolling over on a certain day, add in the number of days until
338            # the next rollover, but offset by 1 since we just calculated the time
339            # until the next day starts.  There are three cases:
340            # Case 1) The day to rollover is today; in this case, do nothing
341            # Case 2) The day to rollover is further in the interval (i.e., today is
342            #         day 2 (Wednesday) and rollover is on day 6 (Sunday).  Days to
343            #         next rollover is simply 6 - 2 - 1, or 3.
344            # Case 3) The day to rollover is behind us in the interval (i.e., today
345            #         is day 5 (Saturday) and rollover is on day 3 (Thursday).
346            #         Days to rollover is 6 - 5 + 3, or 4.  In this case, it's the
347            #         number of days left in the current week (1) plus the number
348            #         of days in the next week until the rollover day (3).
349            # The calculations described in 2) and 3) above need to have a day added.
350            # This is because the above time calculation takes us to midnight on this
351            # day, i.e. the start of the next day.
352            if self.when in RolloverInterval.WEEK_DAYS:
353                day = current_day  # 0 is Monday
354                if day != self.day_of_week:
355                    if day < self.day_of_week:
356                        days_to_wait = self.day_of_week - day
357                    else:
358                        days_to_wait = 6 - day + self.day_of_week + 1
359                    new_rollover_at = result + (
360                        days_to_wait * ONE_DAY_IN_SECONDS
361                    )
362                    if not self.utc:
363                        dst_now = t[-1]
364                        dst_at_rollover = time.localtime(new_rollover_at)[-1]
365                        if dst_now != dst_at_rollover:
366                            if not dst_now:
367                                # DST kicks in before next rollover, so we need to deduct an hour
368                                new_rollover_at -= ONE_HOUR_IN_SECONDS
369                            else:
370                                # DST bows out before next rollover, so we need to add an hour
371                                new_rollover_at += ONE_HOUR_IN_SECONDS
372                    result = new_rollover_at
373        return result
374
375    def should_rollover(self, record: LogRecord) -> bool:
376        """
377        Determine if rollover should occur.
378
379        record is not used, as we are just comparing times, but it is needed so
380        the method signatures are the same
381        """
382        t = int(time.time())
383        if t >= self.rollover_at:
384            return True
385        return False
386
387    async def get_files_to_delete(self) -> List[str]:
388        """
389        Determine the files to delete when rolling over.
390        """
391        dir_name, base_name = os.path.split(self.absolute_file_path)
392        loop = get_running_loop()
393        file_names = await loop.run_in_executor(
394            None, lambda: os.listdir(dir_name)
395        )
396        result = []
397        prefix = base_name + "."
398        plen = len(prefix)
399        for file_name in file_names:
400            if file_name[:plen] == prefix:
401                suffix = file_name[plen:]
402                if self.ext_match.match(suffix):
403                    result.append(os.path.join(dir_name, file_name))
404        if len(result) < self.backup_count:
405            return []
406        else:
407            result.sort(reverse=True)  # os.listdir order is not defined
408            return result[: len(result) - self.backup_count]
409
410    async def _delete_files(self, file_paths: List[str]):
411        loop = get_running_loop()
412        delete_tasks = (
413            loop.run_in_executor(None, lambda: os.unlink(file_path))
414            for file_path in file_paths
415        )
416        await asyncio.gather(*delete_tasks)
417
418    async def do_rollover(self):
419        """
420        do a rollover; in this case, a date/time stamp is appended to the filename
421        when the rollover happens.  However, you want the file to be named for the
422        start of the interval, not the current time.  If there is a backup count,
423        then we have to get a list of matching filenames, sort them and remove
424        the one with the oldest suffix.
425        """
426        if self.stream:
427            await self.stream.close()
428            self.stream = None
429        # get the time that this sequence started at and make it a TimeTuple
430        current_time = int(time.time())
431        dst_now = time.localtime(current_time)[-1]
432        t = self.rollover_at - self.interval
433        if self.utc:
434            time_tuple = time.gmtime(t)
435        else:
436            time_tuple = time.localtime(t)
437            dst_then = time_tuple[-1]
438            if dst_now != dst_then:
439                if dst_now:
440                    addend = ONE_HOUR_IN_SECONDS
441                else:
442                    addend = -ONE_HOUR_IN_SECONDS
443                time_tuple = time.localtime(t + addend)
444        destination_file_path = self.rotation_filename(
445            self.absolute_file_path
446            + "."
447            + time.strftime(self.suffix, time_tuple)
448        )
449        loop = get_running_loop()
450        if await loop.run_in_executor(
451            None, lambda: os.path.exists(destination_file_path)
452        ):
453            await loop.run_in_executor(
454                None, lambda: os.unlink(destination_file_path)
455            )
456        await self.rotate(self.absolute_file_path, destination_file_path)
457        if self.backup_count > 0:
458            files_to_delete = await self.get_files_to_delete()
459            if files_to_delete:
460                await self._delete_files(files_to_delete)
461
462        await self._init_writer()
463        new_rollover_at = self.compute_rollover(current_time)
464        while new_rollover_at <= current_time:
465            new_rollover_at = new_rollover_at + self.interval
466        # If DST changes and midnight or weekly rollover, adjust for this.
467        if (
468            self.when == RolloverInterval.MIDNIGHT
469            or self.when in RolloverInterval.WEEK_DAYS
470        ) and not self.utc:
471            dst_at_rollover = time.localtime(new_rollover_at)[-1]
472            if dst_now != dst_at_rollover:
473                if not dst_now:
474                    # DST kicks in before next rollover, so we need to deduct an hour
475                    addend = -ONE_HOUR_IN_SECONDS
476                else:
477                    # DST bows out before next rollover, so we need to add an hour
478                    addend = ONE_HOUR_IN_SECONDS
479                new_rollover_at += addend
480        self.rollover_at = new_rollover_at
481