<commit_before>(ns puppetlabs.jdbc-util.core
  (:import [com.zaxxer.hikari HikariDataSource]
           [java.util.regex Pattern]
           [org.postgresql.util PGobject PSQLException]
           [org.postgresql.core Utils])
  (:require [clojure.java.jdbc :as jdbc]
            [clojure.string :as str]
            [clojure.tools.logging :as log]
            [puppetlabs.i18n.core :refer [trs trsn]]
            [puppetlabs.kitchensink.core :as ks]))

(defn pg-escape-string
  "Takes an arbitrary string, escapes any SQL quoting meta-characters, and
  returns it in SQL string literal format to be naively interpolated into an SQL
  string. Only for use with unplanned statements (those besides SELECT, INSERT,
  UPDATE, and DELETE), which do not accept query parameters."
  [s]
  (let [escaped (-> (Utils/escapeLiteral nil s true) .toString)]
    (str "'" escaped "'")))

(defn pg-escape-identifier
  "Escape an arbitrary string to make it safe to naively interpolate into an SQL
  string as an identifier. Use of this function is discouraged."
  [s]
  (-> (Utils/escapeIdentifier nil s) .toString))

(defn connection-pool
  "Given a DB spec map containing :subprotocol, :subname, :user, and :password
  keys, return a pooled DB spec map (one containing just the :datasource key
  with a pooled DataSource object as the value). The returned pooled DB spec
  can be passed directly as the first argument to clojure.java.jdbc's
  functions.

  Times out after 30 seconds and throws org.postgresql.util.PSQLException"
  [db-spec]
  (let [ds (doto (HikariDataSource.)
             (.setJdbcUrl (str "jdbc:"
                               (:subprotocol db-spec) ":"
                               (:subname db-spec)))
             (.setUsername (:user db-spec))
             (.setPassword (:password db-spec))
             (.setInitializationFailFast false))]
    {:datasource ds}))

(defmacro with-timeout [timeout-s default & body]
  `(let [f# (future (do ~@body))
         result# (deref f# (* 1000 ~timeout-s) ~default)]
     (future-cancel f#)
     result#))

(def db-status-timeout-secs 4)

(defn db-up?
  [db-spec]
  (let [result (with-timeout db-status-timeout-secs :timeout
                 (try (let [select-42 "SELECT (a - b) AS answer FROM (VALUES ((7 * 7), 7)) AS x(a, b)"
                            [{:keys [answer]}] (jdbc/query db-spec [select-42])]
                        (= answer 42))
                      (catch Exception e
                        (log/warn e (trs "Status check of db failed with error:"))
                        false)))]
    (if (= :timeout result)
      (do (log/warn (trs "Database status check timed out after 4 seconds."))
          false)
      result)))

(defn db-exists?
  [admin-db-spec db-name]
  (-> (jdbc/query admin-db-spec ["SELECT 1 AS exists FROM pg_database WHERE datname = ?" db-name])
    first
    :exists
    (= 1)))

(defn create-db!
  "Given a DB spec that has a user with permission to create databases and that
  connects to a database that isn't `db-name`, the database's name, and the name
  of the user that will own the database, creates the database `db-name` owned by
  `db-owner`, with the DB's encoding set to UTF-8."
  [admin-db-spec db-name db-owner]
  (let [sql (format "CREATE DATABASE %s WITH OWNER %s ENCODING 'UTF8'"
                    (pg-escape-identifier db-name)
                    (pg-escape-identifier db-owner))]
    (jdbc/execute! admin-db-spec [sql] {:transaction? false})))

(defn drop-db!
  "Given a DB spec that has a user with permission to drop the database
  `db-name` and that connects to a database that isn't `db-name` and the safe
  postgres identifier `db-name` (see the docstring for `pg-identifier?` for the
  definition of 'safe'), drop the database named by `db-name`."
  [admin-db-spec db-name]
  (let [sql (format "DROP DATABASE IF EXISTS %s" (pg-escape-identifier db-name))]
    (jdbc/execute! admin-db-spec [sql] {:transaction? false})
    nil))

(defn user-exists?
  [admin-db-spec user-name]
  (-> (jdbc/query admin-db-spec ["SELECT 1 AS present FROM pg_roles WHERE rolname = ?" user-name])
    first
    :present
    (= 1)))

(defn create-user!
  "Given a DB spec that has a user with permission to create users and that
  connects to a database that isn't `db-name`, the safe postgres identifier
  `db-name` (see the docstring of `safe-pg-identifier?` for the definition of
  'safe'), and the safe postgres identifier `db-owner`, create the
  database `db-name` owned by `db-owner`, with the DB's encoding set to UTF-8."
  [admin-db-spec username password]
  (let [sql (format "CREATE ROLE %s WITH LOGIN PASSWORD %s"
                    (pg-escape-identifier username)
                    (pg-escape-string password))]
    (jdbc/execute! admin-db-spec [sql])
    nil))

(defn drop-user!
  [admin-db-spec username]
  (let [sql (format "DROP ROLE IF EXISTS %s" (pg-escape-identifier username))]
    (jdbc/execute! admin-db-spec [sql])
    nil))

(defn public-tables
  "Get the names of all public tables in a database"
  [db-spec]
  (let [query "SELECT table_name FROM information_schema.tables WHERE LOWER(table_schema) = 'public'"]
    (jdbc/query db-spec [query] {:row-fn :table_name})))

(defn drop-public-tables!
  "Drops all public tables in a database. Super dangerous."
  [db-spec]
  (when-let [tables (seq (public-tables db-spec))]
    (jdbc/db-do-commands db-spec (map #(format "DROP TABLE %s CASCADE" %) tables))))

(defn public-functions
  "Get the names of all public functions in a database"
  [db-spec]<commit_msg>(maint) Improve docstrings for user/db fns.<commit_after>(ns puppetlabs.jdbc-util.core
  (:import [com.zaxxer.hikari HikariDataSource]
           [java.util.regex Pattern]
           [org.postgresql.util PGobject PSQLException]
           [org.postgresql.core Utils])
  (:require [clojure.java.jdbc :as jdbc]
            [clojure.string :as str]
            [clojure.tools.logging :as log]
            [puppetlabs.i18n.core :refer [trs trsn]]
            [puppetlabs.kitchensink.core :as ks]))

(defn pg-escape-string
  "Takes an arbitrary string, escapes any SQL quoting meta-characters, and
  returns it in SQL string literal format to be naively interpolated into an SQL
  string. Only for use with unplanned statements (those besides SELECT, INSERT,
  UPDATE, and DELETE), which do not accept query parameters."
  [s]
  (let [escaped (-> (Utils/escapeLiteral nil s true) .toString)]
    (str "'" escaped "'")))

(defn pg-escape-identifier
  "Escape an arbitrary string to make it safe to naively interpolate into an SQL
  string as an identifier. Use of this function is discouraged."
  [s]
  (-> (Utils/escapeIdentifier nil s) .toString))

(defn connection-pool
  "Given a DB spec map containing :subprotocol, :subname, :user, and :password
  keys, return a pooled DB spec map (one containing just the :datasource key
  with a pooled DataSource object as the value). The returned pooled DB spec
  can be passed directly as the first argument to clojure.java.jdbc's
  functions.

  Times out after 30 seconds and throws org.postgresql.util.PSQLException"
  [db-spec]
  (let [ds (doto (HikariDataSource.)
             (.setJdbcUrl (str "jdbc:"
                               (:subprotocol db-spec) ":"
                               (:subname db-spec)))
             (.setUsername (:user db-spec))
             (.setPassword (:password db-spec))
             (.setInitializationFailFast false))]
    {:datasource ds}))

(defmacro with-timeout [timeout-s default & body]
  `(let [f# (future (do ~@body))
         result# (deref f# (* 1000 ~timeout-s) ~default)]
     (future-cancel f#)
     result#))

(def db-status-timeout-secs 4)

(defn db-up?
  [db-spec]
  (let [result (with-timeout db-status-timeout-secs :timeout
                 (try (let [select-42 "SELECT (a - b) AS answer FROM (VALUES ((7 * 7), 7)) AS x(a, b)"
                            [{:keys [answer]}] (jdbc/query db-spec [select-42])]
                        (= answer 42))
                      (catch Exception e
                        (log/warn e (trs "Status check of db failed with error:"))
                        false)))]
    (if (= :timeout result)
      (do (log/warn (trs "Database status check timed out after 4 seconds."))
          false)
      result)))

(defn db-exists?
  "Given a DB spec that connects to a database besides `db-name`, return
  a boolean indicating whether `db-name` exists."
  [admin-db-spec db-name]
  (-> (jdbc/query admin-db-spec ["SELECT 1 AS exists FROM pg_database WHERE datname = ?" db-name])
    first
    :exists
    (= 1)))

(defn create-db!
  "Given a DB spec that has a user with permission to create databases and that
  connects to a database that isn't `db-name`, the database's name, and the name
  of the user that will own the database, creates the database `db-name` owned
  by `db-owner`, with the DB's encoding set to UTF-8."
  [admin-db-spec db-name db-owner]
  (let [sql (format "CREATE DATABASE %s WITH OWNER %s ENCODING 'UTF8'"
                    (pg-escape-identifier db-name)
                    (pg-escape-identifier db-owner))]
    (jdbc/execute! admin-db-spec [sql] {:transaction? false})))

(defn drop-db!
  "Given a DB spec that has a user with permission to drop the database
  `db-name` and that connects to a database that isn't `db-name`, and the
  database's name `db-name`, drop that database named by `db-name`."
  [admin-db-spec db-name]
  (let [sql (format "DROP DATABASE IF EXISTS %s" (pg-escape-identifier db-name))]
    (jdbc/execute! admin-db-spec [sql] {:transaction? false})
    nil))

(defn user-exists?
  "Given a DB spec that connects with a user besides `username`, return
  a boolean indicating whether `username` exists."
  [admin-db-spec username]
  (-> (jdbc/query admin-db-spec ["SELECT 1 AS present FROM pg_roles WHERE rolname = ?" username])
    first
    :present
    (= 1)))

(defn create-user!
  "Given a DB spec that has a user with permission to create roles and isn't
  `username`, create the user `username` with the given `password`. Accepts an
  optional options map where the :superuser? option will make the created user
  a superuser if the value is exactly `true` (not just truthy); note that only
  superusers can create more superusers."
  [admin-db-spec username password]
  (let [sql (format "CREATE ROLE %s WITH LOGIN PASSWORD %s"
                    (pg-escape-identifier username)
                    (pg-escape-string password))]
    (jdbc/execute! admin-db-spec [sql])
    nil))

(defn drop-user!
  "Given a DB spec that connects with a user besides `username` and has
  permisson to drop that user, remove the named user. Note that only superusers
  can remove other superusers; other users can be removed by anyone with the
  CREATEROLE permission."
  [admin-db-spec username]
  (let [sql (format "DROP ROLE IF EXISTS %s" (pg-escape-identifier username))]
    (jdbc/execute! admin-db-spec [sql])
    nil))

(defn public-tables
  "Get the names of all public tables in a database"
  [db-spec]
  (let [query "SELECT table_name FROM information_schema.tables WHERE LOWER(table_schema) = 'public'"]
    (jdbc/query db-spec [query] {:row-fn :table_name})))

(defn drop-public-tables!
  "Drops all public tables in a database. Super dangerous."
  [db-spec]
  (when-let [tables (seq (public-tables db-spec))]
    (jdbc/db-do-commands db-spec (map #(format "DROP TABLE %s CASCADE" %) tables))))

(defn public-functions
  "Get the names of all public functions in a database"
  [db-spec]
{'fraction_non_alphanumeric': 0.1163437309931358, 'fraction_numerical': 0.004431314623338257, 'mean_word_length': 2.808735936465917, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5608852', 'n_tokens_mistral': 3739, 'n_tokens_neox': 3528, 'n_words': 1310}