<commit_before>(ns joplin.elasticsearch.database
  (:use [joplin.core])
  (:require [clojure.string]
            [clojurewerkz.elastisch.native :as es]
            [clojurewerkz.elastisch.native.document :as esd]
            [clojurewerkz.elastisch.native.index :as esi]
            [ragtime.core :refer [Migratable]])
  (:import [org.elasticsearch.action.admin.indices.settings.get GetSettingsRequest]
           [org.elasticsearch.client Client]))

;; ============================================================================
;; ES connection

(def ^:dynamic *es-client* nil)

(defn- ensure-connected []
  (if *es-client* *es-client*
      (throw (Exception. "es-client has not been initialised"))))

(defn- wait-for-ready [shards]
  (-> (ensure-connected)
      .admin
      .cluster
      (.prepareHealth (into-array String []))
      ;; Don't wait for green -- it may never come for the cluster.
      .setWaitForYellowStatus
      (.setWaitForActiveShards shards)
      .execute
      .actionGet))

(defn init [{:keys [cluster host port]}]
  (println "Connecting to ES server" cluster host port)
  (when-not *es-client*
    (println "Building an es client")
    (let [client (es/connect [[host port]] {"cluster.name" cluster})]
      (alter-var-root #'*es-client* (fn [_] client))))
  (println "es client is" *es-client*)
  (println "readiness of es cluster:" (wait-for-ready 0)))

(defn init-local [client]
  (alter-var-root #'*es-client* (fn [_] client)))

;; ============================================================================
;; Handle migration tracking

(def migration-index "migrations")
(def migration-type "migration")
(def migration-document-id "0")

(defn- ensure-migration-index []
  (let [es-client (ensure-connected)]
    (if-not (esi/exists? es-client migration-index)
      (do
        (wait-for-ready 0)
        (esi/create es-client migration-index)
        (wait-for-ready 1)))))

(defn- es-get-applied []
  (let [es-client (ensure-connected)]
    (ensure-migration-index)
    (->> (esd/get es-client migration-index migration-type migration-document-id)
         :_source
         :migrations)))

(defn es-get-applied-migrations []
  (->> (es-get-applied)
       (sort-by second)
       keys
       (map name)))

(defn es-add-migration-id [migration-id]
  (let [es-client (ensure-connected)]
    (esd/put es-client migration-index migration-type migration-document-id
             {:migrations (assoc (es-get-applied) migration-id (java.util.Date.))})))

(defn es-remove-migration-id [migration-id]
  (let [es-client (ensure-connected)]
    (esd/put es-client migration-index migration-type migration-document-id
             {:migrations (dissoc (es-get-applied) (keyword migration-id))})))

;; ============================================================================
;; Data migration

(defn migrate-data
  ([old-index mapping-type new-index]
     (migrate-data old-index mapping-type new-index identity))
  ([old-index mapping-type new-index trans-f]
     (let [es-client (ensure-connected)]
       (dorun
        (->> (esd/search old-index
                         mapping-type
                         :query {:match_all {}})
             (esd/scroll-seq)
             (map :_source)
             (map trans-f)
             (pmap (fn [doc]
                     (esd/create new-index mapping-type doc :id (:_id doc)))))))))

;; ============================================================================
;; Functions for use within migrations

(defn- unwrap-settings [settings]
  (reduce (fn [m s] (assoc m (.key s) (-> s .value .getAsStructuredMap))) {} settings))

(defn get-settings
  "Not implemented for native client in elastisch yet so rolling our own"
  [client]
  (let [req (GetSettingsRequest.)]
    (-> ^Client client .admin .indices (.getSettings req) .actionGet
        .getIndexToSettings unwrap-settings)))

(defn find-index-names [alias-name]
  (->> (get-settings (ensure-connected))
       keys
       (map name)
       (filter #(.startsWith % alias-name))
       sort
       reverse))

(defn assign-alias
  "Add an index to an alias, optionally taking an old index name to be remove from the
   alias"
  [alias-name new-index-name & [old-index-name]]
  (let [ops [{:add {:indices new-index-name :alias alias-name}}]
        ops (if old-index-name
              (conj ops {:remove {:index old-index-name
                                  :aliases alias-name}})
              ops)]
    (println "ops" ops)
    (esi/update-aliases (ensure-connected) ops)))

(defn create-index
  "Create an index with the specified options.  Only the alias name is specified, the
   actual index name is auto-generated to avoid conflicts."
  [alias-name & opts]
  (let [new-index-name (str alias-name "-" (System/currentTimeMillis))
        old-index-name (first (find-index-names alias-name))]
    (apply esi/create (ensure-connected) new-index-name opts)
    (assign-alias alias-name new-index-name old-index-name)
    (wait-for-ready 1)))<commit_msg>returning old and new index names from update-index so that they can be used for data migration<commit_after>(ns joplin.elasticsearch.database
  (:use [joplin.core])
  (:require [clojure.string]
            [clojurewerkz.elastisch.native :as es]
            [clojurewerkz.elastisch.native.document :as esd]
            [clojurewerkz.elastisch.native.index :as esi]
            [ragtime.core :refer [Migratable]])
  (:import [org.elasticsearch.action.admin.indices.settings.get GetSettingsRequest]
           [org.elasticsearch.client Client]))

;; ============================================================================
;; ES connection

(def ^:dynamic *es-client* nil)

(defn- ensure-connected []
  (if *es-client* *es-client*
      (throw (Exception. "es-client has not been initialised"))))

(defn- wait-for-ready [shards]
  (-> (ensure-connected)
      .admin
      .cluster
      (.prepareHealth (into-array String []))
      ;; Don't wait for green -- it may never come for the cluster.
      .setWaitForYellowStatus
      (.setWaitForActiveShards shards)
      .execute
      .actionGet))

(defn init [{:keys [cluster host port]}]
  (println "Connecting to ES server" cluster host port)
  (when-not *es-client*
    (println "Building an es client")
    (let [client (es/connect [[host port]] {"cluster.name" cluster})]
      (alter-var-root #'*es-client* (fn [_] client))))
  (println "es client is" *es-client*)
  (println "readiness of es cluster:" (wait-for-ready 0)))

(defn init-local [client]
  (alter-var-root #'*es-client* (fn [_] client)))

;; ============================================================================
;; Handle migration tracking

(def migration-index "migrations")
(def migration-type "migration")
(def migration-document-id "0")

(defn- ensure-migration-index []
  (let [es-client (ensure-connected)]
    (if-not (esi/exists? es-client migration-index)
      (do
        (wait-for-ready 0)
        (esi/create es-client migration-index)
        (wait-for-ready 1)))))

(defn- es-get-applied []
  (let [es-client (ensure-connected)]
    (ensure-migration-index)
    (->> (esd/get es-client migration-index migration-type migration-document-id)
         :_source
         :migrations)))

(defn es-get-applied-migrations []
  (->> (es-get-applied)
       (sort-by second)
       keys
       (map name)))

(defn es-add-migration-id [migration-id]
  (let [es-client (ensure-connected)]
    (esd/put es-client migration-index migration-type migration-document-id
             {:migrations (assoc (es-get-applied) migration-id (java.util.Date.))})))

(defn es-remove-migration-id [migration-id]
  (let [es-client (ensure-connected)]
    (esd/put es-client migration-index migration-type migration-document-id
             {:migrations (dissoc (es-get-applied) (keyword migration-id))})))

;; ============================================================================
;; Data migration

(defn migrate-data
  ([old-index mapping-type new-index]
     (migrate-data old-index mapping-type new-index identity))
  ([old-index mapping-type new-index trans-f]
     (let [es-client (ensure-connected)]
       (dorun
        (->> (esd/search old-index
                         mapping-type
                         :query {:match_all {}})
             (esd/scroll-seq)
             (map :_source)
             (map trans-f)
             (pmap (fn [doc]
                     (esd/create new-index mapping-type doc :id (:_id doc)))))))))

;; ============================================================================
;; Functions for use within migrations

(defn- unwrap-settings [settings]
  (reduce (fn [m s] (assoc m (.key s) (-> s .value .getAsStructuredMap))) {} settings))

(defn get-settings
  "Not implemented for native client in elastisch yet so rolling our own"
  [client]
  (let [req (GetSettingsRequest.)]
    (-> ^Client client .admin .indices (.getSettings req) .actionGet
        .getIndexToSettings unwrap-settings)))

(defn find-index-names [alias-name]
  (->> (get-settings (ensure-connected))
       keys
       (map name)
       (filter #(.startsWith % alias-name))
       sort
       reverse))

(defn assign-alias
  "Add an index to an alias, optionally taking an old index name to be remove from the
   alias"
  [alias-name new-index-name & [old-index-name]]
  (let [ops [{:add {:indices new-index-name :alias alias-name}}]
        ops (if old-index-name
              (conj ops {:remove {:index old-index-name
                                  :aliases alias-name}})
              ops)]
    (println "ops" ops)
    (esi/update-aliases (ensure-connected) ops)))

(defn create-index
  "Create an index with the specified options.  Only the alias name is specified, the
   actual index name is auto-generated to avoid conflicts. Old and new index names returned
   so user can perform data migration."
  [alias-name & opts]
  (let [new-index-name (str alias-name "-" (System/currentTimeMillis))
        old-index-name (first (find-index-names alias-name))]
    (apply esi/create (ensure-connected) new-index-name opts)
    (assign-alias alias-name new-index-name old-index-name)
    (wait-for-ready 1)
    [old-index-name new-index-name]))
{'fraction_non_alphanumeric': 0.19399175338700178, 'fraction_numerical': 0.0009817396426467701, 'mean_word_length': 2.951512800620636, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 25, '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': '6813617', 'n_tokens_mistral': 3157, 'n_tokens_neox': 3078, 'n_words': 813}