AWS CLI wrapping script by Babashka

aws_logs_core.clj

(ns aws.logs.core
  "AWS CLI wrapper for `logs` service
  https://docs.aws.amazon.com/cli/latest/reference/logs/index.html

  You need to install AWS CLI and run `aws configure` beforehand."
  (:require
   [cheshire.core :as json]
   [clojure.java.shell :as shell]
   [clojure.string :as str]))

(def ^:private default-timeout-sec
  "Timeout sec to wait query completion."
  10)

(defn zoned-date-time->epoch [zdt]
  (..  zdt
       toInstant
       getEpochSecond))

(defn date->epoch-range
  "Convert date to `From`/`To` as Epoch time.
  Return a list like: [From To].
  `From` is 00:00:00 on the specified date.
  `To` is 00:00:00 on the day after the specified date."
  [{:keys [yyyy mm dd zone-id]}]
  (let [zdt (java.time.ZonedDateTime/of
              (Long/parseLong yyyy)
              (Long/parseLong mm)
              (Long/parseLong dd)
              0 0 0 0 zone-id)]
    (->> [zdt (.plusDays zdt 1)]
         (map zoned-date-time->epoch))))

(defn construct-start-query-command
  "Construct a command to start query."
  [{:keys [group-name region queries] :as m}]
  (let [[start-time end-time] (date->epoch-range m)
        extra-query (str/join " | " queries)
        extra-query (cond->> extra-query
                      (seq extra-query) (str " | "))]
    (map str
         (flatten
           ["aws" "logs" "start-query"
            (when region
              ["--region" region])
            "--log-group-name" group-name
            "--start-time" start-time "--end-time" end-time
            "--query-string" (str "fields @timestamp, @message" extra-query " | sort @timestamp desc")]))))

(defn construct-get-query-results-command
  "Construct a command to get query results."
  [region query-id]
  (->> ["aws" "logs" "get-query-results"
        (when region
          ["--region" region])
        "--query-id" query-id]
       flatten
       (remove nil?)))

(defn start-query
  "Start to run query by specified date.
  Argument is the same map as `aws.logs.core/query`."
  [m]
  (let [cmd (construct-start-query-command m)
        res (apply shell/sh cmd)]
    (when (not= 0 (:exit res))
      (throw (ex-info "Failed to start query" res)))
    (-> res :out (json/parse-string true) :queryId)))

(defn describe-completed-queries
  "Fetch completed query information."
  [{:keys [group-name region]}]
  (let [cmd (->> ["aws" "logs" "describe-queries"
                  (when region
                    ["--region" region])
                  "--log-group-name" group-name
                  "--status" "Complete"]
                 flatten
                 (remove nil?))
        res (apply shell/sh cmd)]
    (when (not= 0 (:exit res))
      (throw (ex-info "Failed to describe queies" res)))
    (-> res :out (json/parse-string true) :queries)))

(defn query-completed?
  "Return true if completed queries contain the specified query-id."
  [m query-id]
  (let [completed-ids (->> (describe-completed-queries m) (map :queryId) set)]
    (contains? completed-ids query-id)))

(defn wait-query-completion
  "Wait for query completion.
  ExceptionInfo will be thrown when timed out."
  [m query-id timeout-sec]
  (loop [i 0]
    (when (= timeout-sec i)
      (throw (ex-info "Timeout" {})))
    (when-not (query-completed? m query-id)
      (Thread/sleep 1000)
      (recur (inc i)))))

(defn get-query-results
  "Fetch query results corresponding to the specified query-id.
  `aws.logs.core/wait-query-completion` may be needed
  since this operation doesn't check query completion."
  [region query-id]
  (let [res (apply shell/sh (construct-get-query-results-command region query-id))]
    (when (not= 0 (:exit res))
      (throw (ex-info "Failed to get query results" res)))
    (-> res :out (json/parse-string true) :results
        (->> (map #(:value (nth % 1)))
             (map #(json/parse-string % true))
             (map :message)))))

(defn query
  "Fetch logs by specified date, and return them.

  Argumet is a map like following:
  {:yyyy \"Year with log to be extracted\" ; REQUIRED
   :mm \"Month with log to be extracted\"  ; REQURIED
   :dd \"Day with log to be extracted\"    ; REQUIRED
   :group-name \"Group name for log\"      ; REQUIRED
   :region  \"Region\"                     ; OPTIONAL
   :zone-id `Time zone`                    ; OPTIONAL
   :queries [\"Any additional queries\"]   ; OPTIONAL
   }

  The configured region will be used if `region` is not specified.
  The default `zone-id` is `(java.time.ZoneId/of \"Asia/Tokyo\")`."
  [m]
  (let [m (update m :zone-id #(or % (java.time.ZoneId/of "Asia/Tokyo")))
        timeout (get m :timeout default-timeout-sec)]
    (if-let [query-id (start-query m)]
      (do (wait-query-completion m query-id timeout)
          (get-query-results (:region m) query-id))
      (throw (ex-info "Failed to start query" m)))))