guile-squee/squee.scm

;;; squee --- A guile interface to postgres via the ffi

;; Copyright (C) 2015 Christine Lemmer-Webber <cwebber@dustycloud.org>
;; Copyright (C) 2023 Ludovic Courtès <ludo@gnu.org>

;; This library is free software; you can redistribute it and/or
;; modify it under the terms of the GNU Lesser General Public
;; License as published by the Free Software Foundation; either
;; version 3 of the License, or (at your option) any later version.
;;
;; This library is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;; Lesser General Public License for more details.
;;
;; You should have received a copy of the GNU Lesser General Public
;; License along with this library; if not, write to the Free Software
;; Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

(define-module (squee)
  #:use-module (system foreign)
  #:use-module (rnrs enums)
  #:use-module (ice-9 match)
  #:use-module (ice-9 format)
  #:use-module ((srfi srfi-1) #:select (any))
  #:use-module (srfi srfi-26)
  #:autoload   (ice-9 suspendable-ports) (current-read-waiter)
  #:export (;; The important ones
            connect-to-postgres-paramstring
            exec-query
            pg-conn-finish

            ;; enums and indexes of enums
            conn-status-enum conn-status-enum-index
            polling-status-enum polling-status-index
            exec-status-enum exec-status-enum-index
            transaction-status-enum transaction-status-enum-index
            verbosity-enum verbosity-enum-index
            ping-enum ping-enum-index

            ;; **repl and error messages only!**
            enum-set-ref

            ;; Connection stuff
            <pg-conn> pg-conn? wrap-pg-conn unwrap-pg-conn

            ;; @@: We don't export the result pointer though!
            ;;   as this needs to be cleared to avoid memory
            ;;   leaks...
            ;;
            ;;   We might provide a (exec-with-result-ptr)
            ;;   that cleans up the result pointer after calling
            ;;   some thunk though?
            ;;
            ;;   These are still useful for building your own
            ;;   serializer though...
            result-num-rows result-num-cols result-get-value
            result-serializer-simple-list result-metadata))

(define libpq (dynamic-link "libpq"))

;; ---------------------
;; Enums from libpq-fe.h
;; ---------------------

(define conn-status-enum
  (make-enumeration
   '(connection-ok
     connection-bad
     connection-started connection-made
     connection-awaiting-response connection-auth-ok
     connection-auth-ok connection-setenv
     connection-ssl-startup
     connection-needed)))

(define conn-status-enum-index
  (enum-set-indexer conn-status-enum))

(define polling-status-enum
  (make-enumeration
   '(polling-failed
     polling-reading
     polling-writing
     polling-ok
     polling-active)))

(define polling-status-enum-index
  (enum-set-indexer polling-status-enum))

(define exec-status-enum
  (make-enumeration
   '(empty-query
     command-ok tuples-ok
     copy-out copy-in
     bad-response
     nonfatal-error fatal-error
     copy-both
     single-tuple)))

(define exec-status-enum-index
  (enum-set-indexer exec-status-enum))

(define transaction-status-enum
  (make-enumeration
   '(idle active intrans inerror unknown)))

(define transaction-status-enum-index
  (enum-set-indexer transaction-status-enum))

(define verbosity-enum
  (make-enumeration
   '(terse default verbose)))

(define verbosity-enum-index
  (enum-set-indexer verbosity-enum))

(define ping-enum
  (make-enumeration
   '(ok reject no-response no-attempt)))

(define ping-enum-index
  (enum-set-indexer ping-enum))

(define-wrapped-pointer-type <pg-conn>
  pg-conn?
  wrap-pg-conn unwrap-pg-conn
  (lambda (pg-conn port)
    (format port "#<pg-conn ~x (~a)>"
            (pointer-address (unwrap-pg-conn pg-conn))
            (let ((status (pg-conn-status pg-conn)))
              (cond ((eq? status (conn-status-enum-index 'connection-ok))
                     "connected")
                    ((eq? status (conn-status-enum-index 'connection-bad))
                     (let ((conn-error (pg-conn-error-message pg-conn)))
                       (if (equal? conn-error "")
                           "disconnected"
                           (format #f "disconnected, error: ~s" conn-error))))
                    (#t
                     (symbol->string
                      (pg-conn-status-symbol pg-conn))))))))


;; This one should NOT be exposed to the outside world!  We have our
;; own result structure...

(define-wrapped-pointer-type <result-ptr>
  result-ptr?
  wrap-result-ptr unwrap-result-ptr
  (lambda (result-ptr port)
    (format port "#<result-ptr ~x>"
            (pointer-address (unwrap-result-ptr result-ptr)))))


(define (enum-set-ref enum-set k)
  "Take an ENUM-SET and get the item at position K

This is O(n) but theoretically we don't use it much.
Again, REPL only!"
  (list-ref (enum-set->list enum-set) k))


(define-syntax-rule (define-foreign-libpq name return_type func_name arg_types)
  (define name
    (pointer->procedure return_type
                        (dynamic-func func_name libpq)
                        arg_types)))


(define-foreign-libpq %PQconnectdb '* "PQconnectdb" (list '*))
(define-foreign-libpq %PQstatus int "PQstatus" (list '*))
(define-foreign-libpq %PQerrorMessage '* "PQerrorMessage" (list '*))
(define-foreign-libpq %PQfinish void "PQfinish" (list '*))
(define-foreign-libpq %PQntuples int "PQntuples" (list '*))
(define-foreign-libpq %PQnfields int "PQnfields" (list '*))


;; Synchronous interface.
(define-foreign-libpq %PQexec '* "PQexec" (list '* '*))
(define-foreign-libpq %PQexecParams
  '*  ;; Returns a PGresult
  "PQexecParams"
  (list '* ;; connection
        '* ;; command, a string
        int ;; number of parameters
        '* ;; paramTypes, ok to leave NULL
        '* ;; paramValues, here goes your actual parameters!
        '* ;; paramLengths, ok to leave NULL
        '* ;; paramFormats, ok to leave NULL
        int)) ;; resultFormat... probably 0!

;; Asynchronous interface.
(define-foreign-libpq %PQsocket int "PQsocket" '(*))
(define-foreign-libpq %PQsendQuery int "PQsendQuery" (list '* '*))
(define-foreign-libpq %PQsendQueryParams int "PQsendQueryParams"
  (list '* ;; connection
        '* ;; command, a string
        int ;; number of parameters
        '* ;; paramTypes, ok to leave NULL
        '* ;; paramValues, here goes your actual parameters!
        '* ;; paramLengths, ok to leave NULL
        '* ;; paramFormats, ok to leave NULL
        int))
(define-foreign-libpq %PQconsumeInput int "PQconsumeInput" '(*))
(define-foreign-libpq %PQisBusy int "PQisBusy" '(*))
(define-foreign-libpq %PQgetResult '* "PQgetResult" '(*))


(define-foreign-libpq %PQresultStatus int "PQresultStatus" (list '*))
(define-foreign-libpq %PQresStatus '* "PQresStatus" (list int))
(define-foreign-libpq %PQresultErrorMessage '* "PQresultErrorMessage" (list '*))
(define-foreign-libpq %PQclear void "PQclear" (list '*))

(define-foreign-libpq %PQcmdtuples '* "PQcmdTuples" (list '*))
(define-foreign-libpq %PQntuples int "PQntuples" (list '*))
(define-foreign-libpq %PQnfields int "PQnfields" (list '*))
(define-foreign-libpq %PQgetisnull int "PQgetisnull" (list '* int int))
(define-foreign-libpq %PQgetvalue '* "PQgetvalue" (list '* int int))


;; Via mark_weaver.  Thanks Mark!
;;
;; So, apparently we can use a struct of strings just like an array
;; of strings.  Because magic, and because Mark thinks the C standard
;; allows it enough!

(define (string-pointer-list->string-array ls)
  "Take a list of strings, generate a C-compatible list of free strings"
  (make-c-struct
   (make-list (+ 1 (length ls)) '*)
   (append ls (list %null-pointer))))

(define (pg-conn-status pg-conn)
  "Get the connection status from a postgres connection"
  (%PQstatus (unwrap-pg-conn pg-conn)))

(define (pg-conn-status-symbol pg-conn)
  "Human readable version of the pg-conn status.

Inefficient... don't use this in normal code... it's just for you and
the REPL!  (Well, we do use it for errors, because those are
comparatively \"rare\" so this is okay.)  Compare against the enum
value of the symbol instead."
  (let ((status (pg-conn-status pg-conn)))
    (if (< status (length (enum-set->list conn-status-enum)))
        (enum-set-ref conn-status-enum
                (pg-conn-status pg-conn))
        ;; Weird, this is bigger than our enum of statuses
        (string->symbol
         (format #f "unknown-status-~a" status)))))


(define (pg-conn-error-message pg-conn)
  "Get an error message for this connection"
  (pointer->string (%PQerrorMessage (unwrap-pg-conn pg-conn))))

(define %connection-socket-table
  ;; Map <pg-conn> records to a file port backed by the connection's socket.
  ;; TODO: Avoid this side table.
  (make-weak-key-hash-table))

(define (connection-socket-port pg-conn)          ;internal
  "Return the socket port associated with PG-CONN.  Cache it to avoid
allocating a new one at every call."
  (or (hashq-ref %connection-socket-table pg-conn)
      (let* ((fd (%PQsocket (unwrap-pg-conn pg-conn)))
             (port (fdopen fd "r+0")))
        (set-port-revealed! port 1)               ;closed by libpq
        (hashq-set! %connection-socket-table pg-conn port)
        port)))

(define (pg-conn-finish pg-conn)
  "Close out a database connection.

If the connection is already closed, this simply returns #f."
  (if (eq? (pg-conn-status pg-conn)
           (conn-status-enum-index 'connection-ok))
      (begin
        (%PQfinish (unwrap-pg-conn pg-conn))
        (hashq-remove! %connection-socket-table pg-conn)
        #t)
      #f))

(define (connect-to-postgres-paramstring paramstring)
  "Open a connection to the database via a parameter string"
  (let* ((conn-pointer (%PQconnectdb (string->pointer paramstring)))
         (pg-conn (wrap-pg-conn conn-pointer)))
    ;; 'PQconnectdb' might return a pointer that was previously used for
    ;; another connection, possibly backed by a different file descriptor.
    ;; Thus, remove PG-CONN from the side table.
    (hashq-remove! %connection-socket-table pg-conn)

    (if (eq? conn-pointer %null-pointer)
        (throw 'psql-connect-error
               #f "Unable to establish connection"))
    (let ((status (pg-conn-status pg-conn)))
      (if (eq? status (conn-status-enum-index 'connection-ok))
          pg-conn
          (throw 'psql-connect-error
                (enum-set-ref conn-status-enum status)
                (pg-conn-error-message pg-conn))))))


(define (result-num-rows result-ptr)
  (%PQntuples (unwrap-result-ptr result-ptr)))

(define (result-num-cols result-ptr)
  (%PQnfields (unwrap-result-ptr result-ptr)))

(define (result-get-value result-ptr row col)
  (let ((res (unwrap-result-ptr result-ptr)))
    (and (eqv? (%PQgetisnull res row col) 0)
         (pointer->string
          (%PQgetvalue res row col)))))


;; @@: We ought to also have a vector version...
;;     and other serializations...
(define (result-serializer-simple-list result-ptr)
  "Get a simple list of lists representing the result of the query"
  (let ((rows-range (iota (result-num-rows result-ptr)))
        (cols-range (iota (result-num-cols result-ptr))))
    (map
     (lambda (row-i)
       (map
        (lambda (col-i)
          (result-get-value result-ptr row-i col-i))
        cols-range))
     rows-range)))

;; TODO
(define (result-metadata result-ptr)
  #f)


(define (result-ptr-clear result-ptr)
  (%PQclear (unwrap-result-ptr result-ptr)))

(define (result-error-message result-ptr)
  (%PQresultErrorMessage (unwrap-result-ptr result-ptr)))

(define (wait-for-input pg-conn)
  ((current-read-waiter) (connection-socket-port pg-conn)))

(define (process-result result-ptr serializer)
  "Process the result pointed to by RESULT-PTR, returning a regular value and
data upon success."
  (let ((status (%PQresultStatus result-ptr))
        (result-ptr (wrap-result-ptr result-ptr)))
    (cond
     ;; This is the kind of query that returns tuples
     ((eq? status (exec-status-enum-index 'tuples-ok))
      (let ((serialized-result (serializer result-ptr))
            (metadata (result-metadata result-ptr)))
        ;; Gotta clear the result to prevent memory leaks
        (result-ptr-clear result-ptr)
        (values serialized-result metadata)))

     ;; This doesn't return tuples, eg it's a DELETE or something.
     ((eq? status (exec-status-enum-index 'command-ok))
      (let ((metadata (result-metadata result-ptr))
            (rows (%PQcmdtuples (unwrap-result-ptr result-ptr))))
        ;; Gotta clear the result to prevent memory leaks
        (result-ptr-clear result-ptr)
        ;; Return the number of affected rows.
        (values (string->number
                 (pointer->string rows)) metadata)))

     ;; Uhoh, anything else is an error!
     (#t
      (let ((status-message (pointer->string (%PQresStatus status)))
            (error-message (pointer->string
                            (%PQresultErrorMessage (unwrap-result-ptr
                                                    result-ptr)))))
        (result-ptr-clear result-ptr)
        (throw 'psql-query-error
               ;; @@: Do we need result-status?
               ;; (error-symbol result-status result-error-message)
               (enum-set-ref exec-status-enum status)
               status-message error-message))))))

(define %query-exception
  ;; Cookie to represent an exception thrown.
  (list 'query 'exception))

(define* (exec-query pg-conn command #:optional (params '())
                     #:key (serializer result-serializer-simple-list))
  (let* ((param-pointers
          (map (lambda (param)
                 (if param
                     (string->pointer param)
                     %null-pointer))
               params))
         (command-pointer
          (string->pointer command))
         (param-array-pointer
          (string-pointer-list->string-array param-pointers))
         (conn-pointer (unwrap-pg-conn pg-conn))
         (query-sent?
          (not (zero? (if (null? params)
                          (%PQsendQuery conn-pointer command-pointer)
                          (%PQsendQueryParams conn-pointer command-pointer
                                              (length params)
                                              %null-pointer
                                              param-array-pointer
                                              %null-pointer
                                              %null-pointer 0))))))

    ;; Protect the pointers, and thus the memory regions they point to
    ;; from garbage collection, until %PQexecParams has returned
    (identity param-pointers)
    (identity command-pointer)
    (identity param-array-pointer)

    (unless query-sent?
      (throw 'psql-query-error
             #f #f (pg-conn-error-message pg-conn)))

    ;; Cooperate through the suspendable-port mechanism while waiting for a
    ;; reply.
    (let loop ()
      (wait-for-input pg-conn)

      ;; Consume available input.
      (when (zero? (%PQconsumeInput conn-pointer))
        (throw 'psql-query-error
               #f #f (pg-conn-error-message pg-conn)))

      ;; Is the query done?  If not, try again.
      (unless (zero? (%PQisBusy conn-pointer))
        (loop)))

    ;; Call 'PQgetResult' until it returns NULL.
    (let loop ((result-ptr (%PQgetResult conn-pointer)))
      (call-with-values
          (lambda ()
            (catch 'psql-query-error
              (lambda ()
                (process-result result-ptr serializer))
              (lambda args
                (values %query-exception args))))
        (lambda (value metadata)
          (let ((next-result-ptr (%PQgetResult conn-pointer)))
            (if (null-pointer? next-result-ptr)
                (if (eq? value %query-exception)
                    (apply throw metadata)
                    (values value metadata))
                (loop next-result-ptr))))))))

;; (define conn (connect-to-postgres-paramstring "dbname=sandbox"))