--
-- jsonb <-> Ruby transform.  Functions must opt in with
-- TRANSFORM FOR TYPE jsonb.
--
CREATE EXTENSION jsonb_plruby CASCADE;
NOTICE:  installing required extension "plruby"
-- A jsonb argument arrives as native Ruby data.
CREATE FUNCTION jt_classes(jsonb) RETURNS text
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    walk = lambda { |v|
        case v
        when Hash  then "{#{v.map { |k, e| "#{k}:#{walk.call(e)}" }.join(',')}}"
        when Array then "[#{v.map { |e| walk.call(e) }.join(',')}]"
        else "#{v.class.name}(#{v.inspect})"
        end
    }
    walk.call(args[0])
$$;
SELECT jt_classes('{"s": "text", "i": 5, "f": 2.5, "t": true, "n": null, "a": [1, "two"]}');
                                                 jt_classes                                                  
-------------------------------------------------------------------------------------------------------------
 {a:[Integer(1),String("two")],f:Float(2.5),i:Integer(5),n:NilClass(nil),s:String("text"),t:TrueClass(true)}
(1 row)

-- Raw scalars unwrap.
SELECT jt_classes('42');
 jt_classes  
-------------
 Integer(42)
(1 row)

SELECT jt_classes('"just a string"');
       jt_classes        
-------------------------
 String("just a string")
(1 row)

SELECT jt_classes('null');
  jt_classes   
---------------
 NilClass(nil)
(1 row)

-- Integral numbers arrive as Integer (not Float), big ones included.
CREATE FUNCTION jt_sum(jsonb) RETURNS text
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    total = args[0].sum
    "#{total.class.name}=#{total}"
$$;
SELECT jt_sum('[1, 2, 3]');
  jt_sum   
-----------
 Integer=6
(1 row)

SELECT jt_sum('[9007199254740993, 1]');
          jt_sum          
--------------------------
 Integer=9007199254740994
(1 row)

-- Returning Ruby data into jsonb, no JSON.generate needed.
CREATE FUNCTION jt_build(int) RETURNS jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    {'n' => args[0], 'squares' => (1..args[0]).map { |i| i * i },
     'odd' => args[0].odd?, 'label' => :generated, 'missing' => nil}
$$;
SELECT jt_build(3);
                                      jt_build                                      
------------------------------------------------------------------------------------
 {"n": 3, "odd": true, "label": "generated", "missing": null, "squares": [1, 4, 9]}
(1 row)

SELECT jt_build(4)->'squares'->2 AS third_square;
 third_square 
--------------
 9
(1 row)

-- Round-trip fidelity through both directions.
CREATE FUNCTION jt_echo(jsonb) RETURNS jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    args[0]
$$;
SELECT jt_echo('{"a": [1, 2.5, "x", true, null], "b": {"nested": {"deep": []}}}');
                             jt_echo                             
-----------------------------------------------------------------
 {"a": [1, 2.5, "x", true, null], "b": {"nested": {"deep": []}}}
(1 row)

SELECT jt_echo('[]');
 jt_echo 
---------
 []
(1 row)

SELECT jt_echo('{}');
 jt_echo 
---------
 {}
(1 row)

SELECT jt_echo('3.14');
 jt_echo 
---------
 3.14
(1 row)

-- A rewrite in native Ruby: redact keys anywhere, hashes all the way down.
CREATE FUNCTION jt_redact(doc jsonb, key text) RETURNS jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    redact = lambda { |v|
        case v
        when Hash  then v.to_h { |k, e| [k, k == key ? '[redacted]' : redact.call(e)] }
        when Array then v.map { |e| redact.call(e) }
        else v
        end
    }
    redact.call(doc)
$$;
SELECT jt_redact('{"user": "jd", "password": "s3cret", "kids": [{"password": "also"}]}',
                 'password');
                                   jt_redact                                    
--------------------------------------------------------------------------------
 {"kids": [{"password": "[redacted]"}], "user": "jd", "password": "[redacted]"}
(1 row)

-- A NULL jsonb argument is nil; returning nil is SQL NULL.
CREATE FUNCTION jt_null(jsonb) RETURNS jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    args[0]
$$;
SELECT jt_null(NULL) IS NULL AS null_roundtrip;
 null_roundtrip 
----------------
 t
(1 row)

-- BigDecimal (any Numeric) serializes losslessly through its text form.
CREATE FUNCTION jt_bigdec() RETURNS jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    require 'bigdecimal'
    {'exact' => BigDecimal('0.1') + BigDecimal('0.2')}
$$;
SELECT jt_bigdec();
   jt_bigdec    
----------------
 {"exact": 0.3}
(1 row)

-- Unsupported Ruby objects are rejected cleanly.
CREATE FUNCTION jt_bad() RETURNS jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    {'when' => Time.at(0)}
$$;
SELECT jt_bad();
ERROR:  cannot transform Ruby object of class Time to jsonb
CONTEXT:  PL/Ruby function "jt_bad"
-- Non-finite floats cannot become jsonb numbers.
CREATE FUNCTION jt_inf() RETURNS jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    [Float::INFINITY]
$$;
SELECT jt_inf();
ERROR:  cannot transform infinity to jsonb
CONTEXT:  PL/Ruby function "jt_inf"
-- The transform reaches nested contexts: SETOF rows via return_next...
CREATE FUNCTION jt_batches(jsonb) RETURNS SETOF jsonb
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    args[0].each_slice(2) { |chunk| return_next({'chunk' => chunk}) }
    nil
$$;
SELECT * FROM jt_batches('[1, 2, 3, 4, 5]');
    jt_batches     
-------------------
 {"chunk": [1, 2]}
 {"chunk": [3, 4]}
 {"chunk": [5]}
(3 rows)

-- ...RETURNS TABLE with a jsonb column...
CREATE FUNCTION jt_stats(jsonb) RETURNS TABLE (key text, info jsonb)
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    args[0].each do |k, v|
        return_next({'key' => k, 'info' => {'value' => v, 'class' => v.class.name}})
    end
    nil
$$;
SELECT * FROM jt_stats('{"a": 1, "b": "two"}') ORDER BY key;
 key |                info                 
-----+-------------------------------------
 a   | {"class": "Integer", "value": 1}
 b   | {"class": "String", "value": "two"}
(2 rows)

-- ...composite arguments and returns with jsonb fields...
CREATE TYPE jt_doc AS (id int, body jsonb);
CREATE FUNCTION jt_comp(jt_doc) RETURNS jt_doc
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    doc = args[0]
    {'id' => doc['id'] + 1,
     'body' => doc['body'].merge('seen' => doc['body'].keys.sort)}
$$;
SELECT id, body FROM jt_comp(ROW(1, '{"x": 1, "y": 2}')) AS t;
 id |                 body                 
----+--------------------------------------
  2 | {"x": 1, "y": 2, "seen": ["x", "y"]}
(1 row)

-- ...OUT parameters...
CREATE FUNCTION jt_split(doc jsonb, OUT small jsonb, OUT big jsonb)
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    small = doc.select { |_, v| v.is_a?(Numeric) && v < 10 }
    big   = doc.reject { |_, v| v.is_a?(Numeric) && v < 10 }
$$;
SELECT * FROM jt_split('{"a": 1, "b": 100, "c": 5}');
      small       |    big     
------------------+------------
 {"a": 1, "c": 5} | {"b": 100}
(1 row)

-- ...and jsonb[] arguments and returns, element by element.
CREATE FUNCTION jt_arr(jsonb[]) RETURNS jsonb[]
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    args[0].map { |doc| doc.merge('n' => doc.size) } + [{'extra' => true}]
$$;
SELECT jt_arr(ARRAY['{"a": 1}'::jsonb, '{"b": 2, "c": 3}'::jsonb]);
                                    jt_arr                                     
-------------------------------------------------------------------------------
 {"{\"a\": 1, \"n\": 1}","{\"b\": 2, \"c\": 3, \"n\": 2}","{\"extra\": true}"}
(1 row)

-- Triggers that declare the transform get $_TD jsonb columns as Ruby data,
-- and 'MODIFY' accepts Ruby data back.
CREATE TABLE jt_events (id int, payload jsonb);
CREATE FUNCTION jt_stamp() RETURNS trigger
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    doc = $_TD['new']['payload']
    elog('NOTICE', "payload is a #{doc.class.name} with #{doc.size} key(s)")
    $_TD['new']['payload'] = doc.merge('audited' => true, 'op' => $_TD['event'])
    'MODIFY'
$$;
CREATE TRIGGER jt_events_stamp BEFORE INSERT OR UPDATE ON jt_events
    FOR EACH ROW EXECUTE PROCEDURE jt_stamp();
INSERT INTO jt_events VALUES (1, '{"amount": 42}');
NOTICE:  payload is a Hash with 1 key(s)
UPDATE jt_events SET payload = '{"amount": 99}' WHERE id = 1;
NOTICE:  payload is a Hash with 1 key(s)
SELECT id, payload FROM jt_events;
 id |                     payload                     
----+-------------------------------------------------
  1 | {"op": "UPDATE", "amount": 99, "audited": true}
(1 row)

-- A DELETE trigger reads the old row's jsonb as Ruby data too.
CREATE FUNCTION jt_obituary() RETURNS trigger
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    elog('NOTICE', "deleting amount=#{$_TD['old']['payload']['amount']}")
    nil
$$;
CREATE TRIGGER jt_events_bye BEFORE DELETE ON jt_events
    FOR EACH ROW EXECUTE PROCEDURE jt_obituary();
DELETE FROM jt_events;
NOTICE:  deleting amount=99
-- Without the clause, a trigger sees the String form, as always.
CREATE FUNCTION jt_plain_trig() RETURNS trigger LANGUAGE plruby AS $$
    elog('NOTICE', "payload is a #{$_TD['new']['payload'].class.name}")
    nil
$$;
CREATE TRIGGER jt_events_plain BEFORE INSERT ON jt_events
    FOR EACH ROW EXECUTE PROCEDURE jt_plain_trig();
DROP TRIGGER jt_events_stamp ON jt_events;
INSERT INTO jt_events VALUES (2, '{"x": 1}');
NOTICE:  payload is a String
DROP TABLE jt_events;
-- Without the TRANSFORM clause, jsonb still travels as a String.
CREATE FUNCTION jt_untransformed(jsonb) RETURNS text LANGUAGE plruby AS $$
    args[0].class.name
$$;
SELECT jt_untransformed('{"a": 1}');
 jt_untransformed 
------------------
 String
(1 row)

-- SPI results are not affected by the function's transforms: a jsonb column
-- read through spi_exec is still its String form.
CREATE FUNCTION jt_spi() RETURNS text
TRANSFORM FOR TYPE jsonb
LANGUAGE plruby AS $$
    row = spi_fetch_row(spi_exec(%q{select '{"a": 1}'::jsonb as d}))
    row['d'].class.name
$$;
SELECT jt_spi();
 jt_spi 
--------
 String
(1 row)