|
| 1 | +import itertools |
| 2 | +from datetime import datetime |
| 3 | +from typing import Generator, Iterable, Optional, cast |
| 4 | + |
| 5 | +import typer |
| 6 | +from click import ClickException |
| 7 | +from snowflake.cli._plugins.logs.manager import LogsManager, LogsQueryRow |
| 8 | +from snowflake.cli._plugins.object.commands import NameArgument, ObjectArgument |
| 9 | +from snowflake.cli.api.commands.snow_typer import SnowTyperFactory |
| 10 | +from snowflake.cli.api.identifiers import FQN |
| 11 | +from snowflake.cli.api.output.types import ( |
| 12 | + CommandResult, |
| 13 | + MessageResult, |
| 14 | + StreamResult, |
| 15 | +) |
| 16 | + |
| 17 | +app = SnowTyperFactory() |
| 18 | + |
| 19 | + |
| 20 | +@app.command(name="logs", requires_connection=True) |
| 21 | +def get_logs( |
| 22 | + object_type: str = ObjectArgument, |
| 23 | + object_name: FQN = NameArgument, |
| 24 | + from_: Optional[str] = typer.Option( |
| 25 | + None, |
| 26 | + "--from", |
| 27 | + help="The start time of the logs to retrieve. Accepts all ISO8061 formats", |
| 28 | + ), |
| 29 | + to: Optional[str] = typer.Option( |
| 30 | + None, |
| 31 | + "--to", |
| 32 | + help="The end time of the logs to retrieve. Accepts all ISO8061 formats", |
| 33 | + ), |
| 34 | + refresh_time: int = typer.Option( |
| 35 | + None, |
| 36 | + "--refresh", |
| 37 | + help="If set, the logs will be streamed with the given refresh time in seconds", |
| 38 | + ), |
| 39 | + **options, |
| 40 | +): |
| 41 | + """ |
| 42 | + Retrieves logs for a given object. |
| 43 | + """ |
| 44 | + if refresh_time and to: |
| 45 | + raise ClickException( |
| 46 | + "You cannot set both --refresh and --to parameters. Please check the values" |
| 47 | + ) |
| 48 | + |
| 49 | + from_time = get_datetime_from_string(from_, "--from") if from_ else None |
| 50 | + to_time = get_datetime_from_string(to, "--to") if to else None |
| 51 | + |
| 52 | + if refresh_time: |
| 53 | + logs_stream: Iterable[LogsQueryRow] = LogsManager().stream_logs( |
| 54 | + object_type=object_type, |
| 55 | + object_name=object_name, |
| 56 | + from_time=from_time, |
| 57 | + refresh_time=refresh_time, |
| 58 | + ) |
| 59 | + logs = itertools.chain( |
| 60 | + (MessageResult(log.log_message) for logs in logs_stream for log in logs) |
| 61 | + ) |
| 62 | + else: |
| 63 | + logs_iterable: Iterable[LogsQueryRow] = LogsManager().get_logs( |
| 64 | + object_type=object_type, |
| 65 | + object_name=object_name, |
| 66 | + from_time=from_time, |
| 67 | + to_time=to_time, |
| 68 | + ) |
| 69 | + logs = (MessageResult(log.log_message) for log in logs_iterable) # type: ignore |
| 70 | + |
| 71 | + return StreamResult(cast(Generator[CommandResult, None, None], logs)) |
| 72 | + |
| 73 | + |
| 74 | +def get_datetime_from_string( |
| 75 | + date_str: str, |
| 76 | + name: Optional[str] = None, |
| 77 | +) -> datetime: |
| 78 | + try: |
| 79 | + return datetime.fromisoformat(date_str) |
| 80 | + except ValueError: |
| 81 | + raise ClickException( |
| 82 | + f"Incorrect format for '{name}'. Please use one of approved ISO formats." |
| 83 | + ) |
0 commit comments