brick-sql-exec

// -*- mode: C++; indent-tabs-mode: nil; c-basic-offset: 4 -*-

/*
 * Copyright (C) 2019 Petr Rockai <code@fixp.eu>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#pragma once
#include "brick-assert"
#include "brick-string"
#include "brick-cons"
#include "brick-except"

#include <memory>
#include <vector>
#include <set>
#include <cxxabi.h>
#include <codecvt>
#include <endian.h> /* beNNtoh / htobeNN */
#include <sys/select.h>

#include <libpq-fe.h>

namespace brq::sql
{
    struct error : std::runtime_error
    {
        using std::runtime_error::runtime_error;
    };

    struct txn;

    struct notification
    {
        std::string channel, payload;
        bool valid;
        explicit operator bool() { return valid; }

        notification( std::string ch, std::string p )
            : channel( ch ), payload( p ), valid( true )
        {}

        notification() : valid( false ) {}

        bool operator<( const notification &o ) const
        {
            ASSERT( valid && o.valid );
            if ( channel == o.channel )
                return payload < o.payload;
            else
                return channel < o.channel;
        }
    };

    struct conn
    {
        std::set< notification > _pending;
        std::string _notices;
        PGconn *_handle = nullptr;
        PGconn *handle() { return _handle; }

        const char *errmsg() { return PQerrorMessage( _handle ); }

        conn( std::string s )
        {
            if ( !( _handle = PQconnectdb( s.c_str() ) ) )
                throw std::bad_alloc();
            if ( PQstatus( _handle ) != CONNECTION_OK )
                raise< error >() << "connecting to " << s << " " << errmsg();

            setup_notices();
        }

        void setup_notices()
        {
            auto msg = []( void *_this, const char *msg )
            {
                static_cast< conn * >( _this )->_notices += msg;
            };

            PQsetNoticeProcessor( _handle, msg, this );
        }

        conn( conn &&rhs ) noexcept
            : _pending( std::move( rhs._pending ) ),
              _notices( std::move( rhs._notices ) ),
              _handle( rhs._handle )
        {
            rhs._handle = nullptr;
            setup_notices();
        }

        void exec( std::string_view sql )
        {
            _notices.clear();
            auto result = PQexec( handle(), c_str( sql ) );
            if ( PQresultStatus( result ) != PGRES_COMMAND_OK )
                raise< error >() << "executing " << sql << ": " << errmsg();
        }

        notification check_notify()
        {
            PQconsumeInput( handle() );

            while ( auto notify = PQnotifies( handle() ) )
            {
                _pending.emplace( notify->relname, notify->extra );
                PQfreemem( notify );
            }

            if ( _pending.empty() )
                return {};

            auto rv = *_pending.begin();
            _pending.erase( _pending.begin() );
            return rv;
        }

        notification wait( std::optional< time_t > deadline = std::nullopt )
        {
            if ( auto n = check_notify() )
                return n;

            int sock = PQsocket( handle() );
            fd_set fds;
            FD_ZERO( &fds );
            FD_SET( sock, &fds );

            auto now = ::time( nullptr );
            struct timeval tv = { !deadline || now >= *deadline ? 0 : *deadline - now, 0 };
            DEBUG( "awaiting notifications, timeout =", tv.tv_sec );

            if ( ::select( sock + 1, &fds, nullptr, nullptr, deadline ? &tv : nullptr ) < 0 )
                brq::raise< system_error >() << "select on postgres socket";

            auto n = check_notify();

            if ( !n )
                DEBUG( "check_notify: timeout" );

            return n;
        }

        ~conn()
        {
            if ( _handle )
                PQfinish( _handle );
        }
    };

    template< typename col > using get_type = typename col::type;

    template< typename cols >
    struct row
    {
        using columns_t = cols;
        using values_t  = typename cols::template map_t< get_type >;
        values_t _values;

        row() = default;
        row( const cols &v ): _values( v ) {}

        template< typename col, typename cs_t, typename vs_t >
        auto &get( vs_t &vs )
        {
            if constexpr ( std::is_same_v< col, typename cs_t::car_t > )
                return vs.car();
            else
                return get< col, typename cs_t::cdr_t >( vs.cdr() );
        }

        template< int i > auto &get() { return get< i >( _values ); }
        template< int i, typename vs_t > auto &get( vs_t &vs )
        {
            if constexpr ( i == 0 )
                return vs.car();
            else
                return get< i - 1, typename vs_t::cdr_t >( vs.cdr() );
        }

        template< typename col >
        auto &get() { return get< col, columns_t >( _values ); }
    };

    template< typename... cols >
    struct iterator : with_columns< cols... >
    {
        using columns_t = typename with_columns< cols... >::columns;

        PGresult *_result = nullptr;
        int _row = 0, _col = 0;

        template< typename rhs_t >
        std::enable_if_t< std::is_same_v< columns_t, typename rhs_t::columns_t >, iterator & >
        operator=( const rhs_t &rhs )
        {
            _result = rhs._result;
            _row = rhs._row;
            _col = rhs._cols;
        }

        bool operator==( const iterator &o ) const { return o._result == _result && o._row == _row; }
        bool operator!=( const iterator &o ) const { return !( *this == o ); }

        iterator() = default;
        iterator( PGresult *res, int row ) : _result( res ), _row( row ) {}
        iterator &operator++() { ++_row; _col = 0; return *this; }
        iterator operator++( int ) { auto r = *this; ++ (*this); return r; }

        template< typename T >
        void fetch_raw( T &t )
        {
            ASSERT( !PQgetisnull( _result, _row, _col ), "row =", _row, "col =", _col );
            ASSERT_EQ( PQgetlength( _result, _row, _col ), sizeof( T ) );
            std::memcpy( &t, PQgetvalue( _result, _row, _col ), sizeof( T ) );
        }

        void fetch_one( inet &in ) { fetch_raw( in ); }
        void fetch_one( bool &t ) { fetch_raw( t ); }
        void fetch_one( int16_t &t ) { fetch_raw( t ); t = be16toh( t ); }
        void fetch_one( int32_t &t ) { fetch_raw( t ); t = be32toh( t ); }
        void fetch_one( int64_t &t ) { fetch_raw( t ); t = be64toh( t ); }
        void fetch_one( timestamp &t ) { fetch_raw( t.fract ); t.fract = be64toh( t.fract ); }
        template< typename T, int s > void fetch_one( fixnum< s, T > &n ) { fetch_one( n.value ); }

        template< typename T, size_t s > void fetch_one( std::array< T, s > &n )
        {
            ASSERT_EQ( PQgetlength( _result, _row, _col ), n.size() * sizeof( T ) );
            std::memcpy( n.begin(), PQgetvalue( _result, _row, _col ), n.size() * sizeof( T ) );
        }

        void fetch_one( interval &i )
        {
            fetch_raw( i );
            i.seconds = be64toh( i.seconds );
            i.days    = be32toh( i.days );
            i.months  = be32toh( i.months );
        }

        void fetch_one( std::u32string &us )
        {
            const char *data = PQgetvalue( _result, _row, _col );
            std::wstring_convert< std::codecvt_utf8< char32_t >, char32_t > conv;
            us = conv.from_bytes( data, data + PQgetlength( _result, _row, _col ) );
        }

        void fetch_one( std::string &s )
        {
            const char *data = PQgetvalue( _result, _row, _col );
            s = std::string( data, data + PQgetlength( _result, _row, _col ) );
        }

        template< typename T > void fetch_one( std::optional< T > &opt )
        {
            if ( PQgetisnull( _result, _row, _col ) )
                opt = std::nullopt;
            else
                opt = T(), fetch_one( *opt );
        }

        row< columns_t > operator*()
        {
            row< columns_t > r;
            r._values.each( [&]( auto &v ) { fetch_one( v ); ++ _col; } );
            return r;
        }
    };

    template< typename self_t >
    struct binder
    {
        self_t &self() { return *static_cast< self_t * >( this ); }

        template< typename T > void bind_raw( const T &t )
        {
            bind_mem( reinterpret_cast< const char * >( &t ), sizeof( T ) );
        }

        void bind_one( interval &i )
        {
            i.seconds = htobe64( i.seconds );
            i.days = htobe32( i.days );
            i.months = htobe32( i.months );
            bind_raw( i );
        }

        void bind_one( timestamp &t ) { t.fract = htobe64( t.fract ); bind_raw( t ); }
        void bind_one( date &d ) { d.julian -= 2451545; d.julian = htobe32( d.julian ); bind_raw( d ); }
        void bind_one( int64_t &t ) { t = htobe64( t ); bind_raw( t ); }
        void bind_one( int32_t &t ) { t = htobe32( t ); bind_raw( t ); }
        void bind_one( int16_t &t ) { t = htobe16( t ); bind_raw( t ); }
        void bind_one( bool &t )    { bind_raw( t ); }
        void bind_one( inet &in )   { bind_raw( in ); }
        void bind_one( std::string_view s ) { bind_mem( s.data(), s.size() ); }
        void bind_one( std::u32string_view us )
        {
            std::wstring_convert< std::codecvt_utf8< char32_t >, char32_t > conv;
            auto str = conv.to_bytes( &*us.begin(), &*us.end() );
            bind_mem( str.data(), str.size() );
        }

        template< typename T, int s > void bind_one( fixnum< s, T > &f ) { bind_one( f.value ); }

        template< size_t n >
        void bind_one( const std::array< uint8_t, n > &s ) { bind_mem( s.begin(), s.size() ); }

        template< typename T >
        std::enable_if_t< T::is_enum_val > bind_one( T v )
        {
            bind_one( std::string( get_name( v, 0 ) ) );
        }

        template< typename T >
        std::enable_if_t< T::is_enum && ! T::is_enum_val > bind_one( T v )
        {
            v._value.match( [&]( auto x ) { bind_one( std::string( get_name( x, 0 ) ) ); } );
        }

        template< typename T > void bind_one( std::optional< T > v )
        {
            if ( v.has_value() )
                bind_one( *v );
            else
                self().bind_null();
        }

        void bind_mem( const uint8_t *data, int size )
        {
            bind_mem( reinterpret_cast< const char * >( data ), size );
        }

        void bind_mem( const char*data, int size ) { self().bind_mem( data, size ); }
    };

    struct copy_base : binder< copy_base >
    {
        conn &_conn;
        bool _prepend_size = false;
        bool _open = true;

        void bind_mem( const char *data, int size )
        {
            if ( _prepend_size )
                _prepend_size = false, put32( size );
            PQputCopyData( _conn.handle(), data, size );
        }

        void bind_null() { put32( -1 ); }

        void put32( int32_t v ) { this->bind_one( v ); }
        void put16( int16_t v ) { this->bind_one( v ); }

        void close()
        {
            if ( !_open ) return;
            _open = false;
            put16( -1 );
            PQputCopyEnd( _conn.handle(), nullptr );
            auto r = PQgetResult( _conn.handle() );
            auto s = PQresultStatus( r );
            if ( s != PGRES_COMMAND_OK )
                raise< error >() << "executing copy: " << _conn.errmsg();
        }

        copy_base( conn &c ) : _conn( c )
        {
            this->bind_mem( "PGCOPY\n\377\r\n\0", 11 );
            put32( 0 );
            put32( 0 );
        }

        ~copy_base() { close(); }
    };

    template< typename cols >
    struct copy_in : copy_base
    {
        using columns = cols;
        using values = typename columns::template map_t< get_type >;

        using copy_base::copy_base;

        copy_in( const copy_in & ) = delete;
        template< typename... T > void put( const T &... args ) { put_list( values( args... ) ); }

        void put_list( values list )
        {
            put16( values::size );
            list.each( [&]( auto &v ) { _prepend_size = true; this->bind_one( v ); } );
        }
    };

    struct stmt_base : binder< stmt_base >
    {
        struct
        {
            sql::conn *conn = nullptr;
            brq::string_builder query, debug;

            std::vector< std::unique_ptr< char[] > > memory;
            std::vector< const char * > params;
            std::vector< int > lengths;
            PGresult *result = nullptr;

            auto reset() { auto r = std::move( *this ); *this = {}; return r; }
        } _d;

        int size() const { return PQntuples( _d.result ); }

        void bind_mem( const char *data, int size )
        {
            _d.lengths.push_back( size );
            _d.memory.emplace_back( new char[ size ] );
            std::memcpy( _d.memory.back().get(), data, size );
            _d.params.push_back( _d.memory.back().get() );
        }

        void bind_null()
        {
            _d.lengths.push_back( 0 );
            _d.params.push_back( nullptr );
        }

        template< typename T >
        stmt_base( conn &c, const T &t )
        {
            _d.conn = &c;
            _d.query << t;
            _d.debug << "where";
        }

        stmt_base() = default;
        stmt_base( const stmt_base & ) = delete;
        stmt_base( stmt_base &&o ) : _d( o._d.reset() ) {}
        stmt_base &operator=( stmt_base &&o ) { _d = o._d.reset(); return *this; }

        ~stmt_base()
        {
            if ( _d.result )
                PQclear( _d.result );
        }

        void exec()
        {
            std::vector< int > formats( _d.params.size(), 1 );

            DEBUG( _d.query.data(), _d.params.size() ? _d.debug.data() : "" );
            _d.conn->_notices.clear();
            _d.result = PQexecParams( _d.conn->handle(), brq::c_str( _d.query.data() ),
                                      _d.params.size(), nullptr, _d.params.data(),
                                      _d.lengths.data(), formats.data(), 1 );

            auto r = PQresultStatus( _d.result );
            if ( r != PGRES_COMMAND_OK && r != PGRES_TUPLES_OK && r != PGRES_COPY_IN )
                raise< error >() << "executing " << _d.query.data() << ": " << _d.conn->errmsg();
        }

        template< typename list_t >
        void bind_list( list_t list )
        {
            int count = 0;
            list.each( [&]( auto ) { ++count; } );
            _d.params.reserve( _d.params.size() + count );
            _d.memory.reserve( _d.memory.size() + count );

            int i = _d.params.size();
            list.each( [&]( auto &v ) { _d.debug << " $" << ++i << " = '" << v << "'"; } );
            list.each( [&]( auto &v ) { bind_one( v ); } );
        }

        template< typename... args >
        void bind( const args &... vs )
        {
            bind_list( cons_list( vs... ) );
        }
    };

    template< typename... cols >
    struct stmt : stmt_base, with_columns< cols... >
    {
        using stmt_base::stmt_base;
        using columns_t = typename with_columns< cols... >::columns;
        using result_t = row< columns_t >;

        using iterator = sql::iterator< cols... >;
        iterator begin() { return iterator( _d.result, 0 ); }
        iterator end()   { return iterator( _d.result, size() ); }

        template< typename rhs_t >
        std::enable_if_t< std::is_same_v< columns_t, typename rhs_t::columns_t >, stmt & >
        operator=( const rhs_t &rhs )
        {
            return *this;
        }

        row< columns_t > result()
        {
            ASSERT_EQ( PQntuples( _d.result ), 1 );
            return *begin();
        }

        stmt &exec()
        {
            stmt_base::exec();
            return *this;
        }

        template< typename... args >
        stmt &bind( const args &... vs )
        {
            stmt_base::bind( vs... );
            return *this;
        }
    };

    struct txn_base
    {
        enum isolation { read_committed, repeatable_read, serializable };

        sql::conn *_conn = nullptr;
        bool _closed = true;
        bool _writable = true;
        isolation _isolation = read_committed;

        explicit txn_base( sql::conn &conn ) : _conn( &conn ) {}
        txn_base() = default;

        void set_isolation( isolation i )
        {
            ASSERT( _closed );
            _isolation = i;
        }

        void set_writable( bool w )
        {
            ASSERT( _closed );
            _writable = w;
        }

        virtual bool open()
        {
            if ( _closed )
            {
                _closed = false;
                exec( "begin transaction" );
                exec() << "set transaction isolation level "
                       << ( _isolation == read_committed  ? "read committed"  :
                            _isolation == repeatable_read ? "repeatable read" :
                            _isolation == serializable    ? "serializable" : "error" )
                       << " "
                       << ( _writable ? "read write" : "read only" );
                return true;
            }

            return false;
        }

        bool valid() const
        {
            return _conn && !_closed;
        }

        txn_base( txn_base&& rhs ) noexcept
            : _conn( rhs._conn ), _closed( rhs._closed )
        {
            rhs._closed = true;
        }

        txn_base& operator=( txn_base &&rhs )
        {
            rollback();
            _conn = rhs._conn;
            _closed = rhs._closed;
            rhs._closed = true;
            return *this;
        }

        txn_base& operator=( const txn_base &rhs ) = delete;

        sql::conn &conn() { return *_conn; }

        void commit()
        {
            if ( !_closed )
            {
                exec( "commit transaction" );
                _closed = true;
            }
        }

        void rollback() noexcept
        {
            if ( !_closed )
            {
                exec( "rollback transaction" );
                _closed = true;
            }
        }

        struct savepoint
        {
            txn_base *_txn;
            std::string _name;

            savepoint( txn_base *t, std::string_view get_name ) : _txn( t ), _name( get_name )
            {
                _txn->exec( brq::format( "savepoint ", get_name ).data() );
            }

            ~savepoint()
            {
                DEBUG( "~savepoint", std::uncaught_exceptions() );

                if ( !_txn )
                    return;

                if ( std::uncaught_exceptions() )
                    rollback();
                else
                    release();
            }

            void rollback()
            {
                _txn->exec( brq::format( "rollback to savepoint ", _name ).data() );
                _txn = nullptr;
            }

            void release()
            {
                _txn->exec( brq::format( "release savepoint ", _name ).data() );
                _txn = nullptr;
            }
        };

        savepoint save( std::string_view name = "save_except" )
        {
            return { this, name };
        }

        void exec( std::string_view q )
        {
            open();
            stmt<> s( conn(), q );
            s.exec();
        }

        template< typename... v >
        void exec( std::string_view q, v... vals )
        {
            open();
            stmt<> s( conn(), q );
            s.bind( vals... );
            s.exec();
        }

        template< typename query_t >
        std::enable_if_t< query_t::is_query, typename query_t::stmt > exec( query_t q )
        {
            return q.template _exec_query< typename query_t::stmt >( conn(), q );
        }

        template< typename tab >
        sql::copy_in< typename tab::columns > copy_in()
        {
            string_builder b;
            b <<  "COPY " << tab::get_name( tab() ) << " ( ";
            using cols = typename tab::columns;
            int i = 0;
            cols().each( [&]( const auto &c ) { b << ( i++ ? ", " : "" ) << c.get_name( c ); } );
            b << " ) FROM STDIN BINARY";
            exec( b.data() );
            return { conn() };
        }

        struct autoexec : brq::string_builder
        {
            txn_base &_txn;
            autoexec( autoexec && ) = default;
            autoexec( txn_base &t ) : _txn( t ) {}
            ~autoexec() { if ( _d.buffer ) _txn.exec( data() ); }

            template< typename T >
            autoexec &operator<<( const T &x )
            {
                *static_cast< brq::string_builder * >( this ) << x;
                return *this;
            }
        };

        autoexec exec() { return { *this }; }

        template< typename t >
        std::enable_if_t< t::is_table > notify() { notify( t::get_name( t() ) ); }

        template< typename t >
        std::enable_if_t< t::is_table > listen() { exec() << "listen " << t::get_name( t() ); }

        template< typename t >
        std::enable_if_t< t::is_table || t::is_view > lock( std::string_view mode )
        {
            exec() << "lock " << t::get_name( t() ) << " in " << mode << " mode";
        }

        void notify( std::string_view ch, std::string_view p = "" )
        {
            exec() << "notify " << ch << ", '" << p << "'";
        }

        void listen( std::string_view sv ) { exec() << "listen " << sv; }

        notification wait( std::optional< time_t > deadline = std::nullopt )
        {
            ASSERT( _closed ); /* waiting with an open transaction causes a deadlock */
            return conn().wait( deadline );
        }

        template< typename... ts >
        std::enable_if_t<( ts::is_table && ... )> truncate()
        {
            brq::string_builder b;
            b << "truncate ";
            int count = 0;
            auto print_one = [&]( auto &t )
            {
                b << ( count++ ? ", " : "" ) << t.get_name( t );
            };
            brq::cons_list_t< ts... >().each( print_one );
            exec( b.data() );
        }

        ~txn_base() noexcept { if ( !_closed ) rollback(); }
    };
}

namespace brq
{
    using sql_connection  = sql::conn;
    using sql_transaction = sql::txn;
    using sql_error       = sql::error;
    template< typename... args >
    using sql_statement   = sql::stmt< args... >;
}

namespace std
{
    template< typename cols >
    struct tuple_size< brq::sql::row< cols > >
        : std::integral_constant< std::size_t, cols::size >
    {};

    template< size_t I, typename cols >
    struct tuple_element< I, brq::sql::row< cols > >
    {
        using type = typename cols::template type_at< I >::type;
    };
}