<commit_before>(ns shale.redis
  (:require [taoensso.carmine :as car :refer (wcar)]
            [schema.core :as s]
            [shale.utils :refer :all]))

(defn hset-all [k m]
  (doseq [[a b] m]
    (car/hset k a b)))

(defn sset-all [k s]
  (car/del k)
  (doseq [a s]
    (car/sadd k a)))

;; in-database models
(def redis-key-prefix "_shale")

(def session-set-key
  (apply str (interpose "/" [redis-key-prefix "sessions"])))

(def session-key-template
  (apply str (interpose "/" [redis-key-prefix "sessions" "%s"])))

(def session-tags-key-template
  (apply str (interpose "/" [redis-key-prefix "sessions" "%s" "tags"])))

(defn session-key [session-id]
  (format session-key-template session-id))

(defn session-tags-key [session-id]
  (format session-tags-key-template session-id))

(def node-set-key
  (apply str (interpose "/" [redis-key-prefix "nodes"])))

(def node-key-template
  (apply str (interpose "/" [redis-key-prefix "nodes" "%s"])))

(def node-tags-key-template
  (apply str (interpose "/" [redis-key-prefix "nodes" "%s" "tags"])))

(defn node-key [id]
  (format node-key-template id))

(defn node-tags-key [id]
  (format node-tags-key-template id))

(def soft-delete-sub-key
  (apply str (interpose "/" [redis-key-prefix "deleting"])))

;; model schemas
(defmacro defmodel
  "A wrapper around schema's `defschema` to attach additional model metadata.

  For example,
  ```
  > (defmodel Model
      \"My Redis model.\"
      {(s/optional-key :name) s/Str})
  > (meta Model)
  {:doc \"My Redis model.\"
   :name \"Model\"
   :model-name \"Model\"
   :redis-key \"_shale/Model\"}

  > (defmodel FancyModel
      \"My fancy model.\"
      :model-name \"fancy\"
      {(s/optional-key :field) s/Int})
  > (meta FancyModel)
  {:doc \"My fancy model.\"
   :name \"Fancy\"
   :model-name \"fancy\"
   :redis-key \"_shale/fancy\"}
  ```"
  [name & forms]
  (let [optional-args (butlast forms)
        docstring (if (string? (first optional-args))
                    (first optional-args))
        options (s/validate (s/either (s/eq nil)
                                      (s/pred #(even? (count %))))
                            (if (nil? docstring)
                              optional-args
                              (rest optional-args)))
        schema (last forms)
        option-map (->> options
                        (partition 2)
                        (map vec)
                        (into {}))
        meta-options (update-in option-map
                                [:model-name]
                                #(or % name))
        default-redis-key (->> [redis-key-prefix (:model-name meta-options)]
                               (clojure.string/join "/"))]
    `(def ~name
       ~(or docstring "")
       (merge-meta (schema.core/schema-with-name ~schema '~name)
                   (merge {:redis-key ~default-redis-key}
                          ~meta-options)))))

(defmodel SessionInRedis
  "A session, as represented in redis."
  :model-name "sessions"
  {(s/optional-key :webdriver-id)   s/Str
   (s/optional-key :tags)         #{s/Str}
   (s/optional-key :reserved)       s/Bool
   (s/optional-key :current-url)    s/Str
   (s/optional-key :browser-name)   s/Str
   (s/optional-key :node-id)        s/Str})

(defmodel NodeInRedis
  "A node, as represented in redis."
  :model-name "nodes"
  {(s/optional-key :url)            s/Str
   (s/optional-key :max-sessions)   s/Int
   (s/optional-key :tags)         #{s/Str}})

(s/defschema IPAddress
  (s/pred #(re-matches #"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}" %)))

(defmodel ProxyInRedis
  "A proxy, as represented in redis."
  :model-name "proxies"
  {(s/optional-key :public-ip) IPAddress
   :type                       (s/enum :socks5 :http)
   :private-host-and-port      s/Str
   (s/optional-key :active)    s/Bool})

;; model fetching
(defn model-key [model-schema id]
  (clojure.string/join "/" [(:redis-key (meta model-schema)) id]))
<commit_msg>Use the is-ip? util for proxies.<commit_after>(ns shale.redis
  (:require [taoensso.carmine :as car :refer (wcar)]
            [schema.core :as s]
            [shale.utils :refer :all]))

(defn hset-all [k m]
  (doseq [[a b] m]
    (car/hset k a b)))

(defn sset-all [k s]
  (car/del k)
  (doseq [a s]
    (car/sadd k a)))

;; in-database models
(def redis-key-prefix "_shale")

(def session-set-key
  (apply str (interpose "/" [redis-key-prefix "sessions"])))

(def session-key-template
  (apply str (interpose "/" [redis-key-prefix "sessions" "%s"])))

(def session-tags-key-template
  (apply str (interpose "/" [redis-key-prefix "sessions" "%s" "tags"])))

(defn session-key [session-id]
  (format session-key-template session-id))

(defn session-tags-key [session-id]
  (format session-tags-key-template session-id))

(def node-set-key
  (apply str (interpose "/" [redis-key-prefix "nodes"])))

(def node-key-template
  (apply str (interpose "/" [redis-key-prefix "nodes" "%s"])))

(def node-tags-key-template
  (apply str (interpose "/" [redis-key-prefix "nodes" "%s" "tags"])))

(defn node-key [id]
  (format node-key-template id))

(defn node-tags-key [id]
  (format node-tags-key-template id))

(def soft-delete-sub-key
  (apply str (interpose "/" [redis-key-prefix "deleting"])))

;; model schemas
(defmacro defmodel
  "A wrapper around schema's `defschema` to attach additional model metadata.

  For example,
  ```
  > (defmodel Model
      \"My Redis model.\"
      {(s/optional-key :name) s/Str})
  > (meta Model)
  {:doc \"My Redis model.\"
   :name \"Model\"
   :model-name \"Model\"
   :redis-key \"_shale/Model\"}

  > (defmodel FancyModel
      \"My fancy model.\"
      :model-name \"fancy\"
      {(s/optional-key :field) s/Int})
  > (meta FancyModel)
  {:doc \"My fancy model.\"
   :name \"Fancy\"
   :model-name \"fancy\"
   :redis-key \"_shale/fancy\"}
  ```"
  [name & forms]
  (let [optional-args (butlast forms)
        docstring (if (string? (first optional-args))
                    (first optional-args))
        options (s/validate (s/either (s/eq nil)
                                      (s/pred #(even? (count %))))
                            (if (nil? docstring)
                              optional-args
                              (rest optional-args)))
        schema (last forms)
        option-map (->> options
                        (partition 2)
                        (map vec)
                        (into {}))
        meta-options (update-in option-map
                                [:model-name]
                                #(or % name))
        default-redis-key (->> [redis-key-prefix (:model-name meta-options)]
                               (clojure.string/join "/"))]
    `(def ~name
       ~(or docstring "")
       (merge-meta (schema.core/schema-with-name ~schema '~name)
                   (merge {:redis-key ~default-redis-key}
                          ~meta-options)))))

(defmodel SessionInRedis
  "A session, as represented in redis."
  :model-name "sessions"
  {(s/optional-key :webdriver-id)   s/Str
   (s/optional-key :tags)         #{s/Str}
   (s/optional-key :reserved)       s/Bool
   (s/optional-key :current-url)    s/Str
   (s/optional-key :browser-name)   s/Str
   (s/optional-key :node-id)        s/Str})

(defmodel NodeInRedis
  "A node, as represented in redis."
  :model-name "nodes"
  {(s/optional-key :url)            s/Str
   (s/optional-key :max-sessions)   s/Int
   (s/optional-key :tags)         #{s/Str}})

(s/defschema IPAddress
  (s/pred is-ip?))

(defmodel ProxyInRedis
  "A proxy, as represented in redis."
  :model-name "proxies"
  {(s/optional-key :public-ip) IPAddress
   :type                       (s/enum :socks5 :http)
   :private-host-and-port      s/Str
   (s/optional-key :active)    s/Bool})

;; model fetching
(defn model-key [model-schema id]
  (clojure.string/join "/" [(:redis-key (meta model-schema)) id]))

{'fraction_non_alphanumeric': 0.16823260538344337, 'fraction_numerical': 0.0015236160487557136, 'mean_word_length': 2.093872741555381, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4291082', 'n_tokens_mistral': 2806, 'n_tokens_neox': 2748, 'n_words': 654}