Commit 02076be2 authored by Kirill Smelkov's avatar Kirill Smelkov

go/*: Move creation of a Pickler and Unpickler to central place

We already have 3 places where we create picklers and unpicklers:
zodb/pydata.go, fs1/index.go and zeo/proto.go and they are already
diverging a bit: for example pydata was explicitly specifying protocol
for pickle encoder, while other places were using defaults.

We will soon want to use StrictUnicode option for all picklers and
unpicklers we create, which means it is better to keep this
customization only in one place instead of copy-pasting it in between
callsites increasing the possibility for divergence and errors.

-> Move the code to create pickler and unpickler to internal zodbpickle
package as a preparatory step for that.

For the reference: this package is merely a wrapper around ogórek, not a
fork of any code like zodbpickle/py is.
parent 8bc46475
// Copyright (C) 2016-2024 Nexedi SA and Contributors.
// Kirill Smelkov <kirr@nexedi.com>
//
// This program is free software: you can Use, Study, Modify and Redistribute
// it under the terms of the GNU General Public License version 3, or (at your
// option) any later version, as published by the Free Software Foundation.
//
// You can also Link and Combine this program with other software covered by
// the terms of any of the Free Software licenses or any of the Open Source
// Initiative approved licenses and Convey the resulting work. Corresponding
// source of such a combination shall include the source code for all other
// software used.
//
// This program is distributed WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//
// See COPYING file for full licensing terms.
// See https://www.nexedi.com/licensing for rationale and options.
// Package zodbpickle provides pickling facility tailored to ZODB/go needs.
//
// Contrary to zodbpickle/py it merely wraps [github.com/kisielk/og-rek] without
// forking code from anywhere.
//
// Use [NewPickler]/[NewUnpickler] to create pickle encoder and decoder
// correspondingly.
//
// The encoder emits pickles with protocol=2 in order to support pristine python2.
//
// See package [github.com/kisielk/og-rek] for details of pickling/unpickling on Go side.
package zodbpickle
import (
"io"
pickle "github.com/kisielk/og-rek"
)
// NewPickler creates new pickle encoder.
//
// The encoder adheres to semantic explained in top-level documentation of
// zodbpickle package.
//
// getref, if !nil, will be used by the encoder to see if a referenced object
// should be represented by persistent reference in generated pickle stream.
// See documentation of ogórek.EncoderConfig.PersistentRef for details.
func NewPickler(w io.Writer, getref func(obj any) *pickle.Ref) *pickle.Encoder {
return pickle.NewEncoderWithConfig(w, &pickle.EncoderConfig{
// allow pristine python2 to decode the pickle.
// TODO 2 -> 3 since ZODB5 switched to it and uses zodbpickle.
Protocol: 2, // see top-level doc
PersistentRef: getref,
})
}
// NewUnpickler creates new pickle decoder.
//
// loadref, if !nil, will be used by the decoder to handle persistent
// references. See documentation of ogórek.DecoderConfig.PersistentLoad for details.
func NewUnpickler(r io.Reader, loadref func(ref pickle.Ref) (any, error)) *pickle.Decoder {
return pickle.NewDecoderWithConfig(r, &pickle.DecoderConfig{
PersistentLoad: loadref,
})
}
// Copyright (C) 2016-2019 Nexedi SA and Contributors. // Copyright (C) 2016-2024 Nexedi SA and Contributors.
// Kirill Smelkov <kirr@nexedi.com> // Kirill Smelkov <kirr@nexedi.com>
// //
// This program is free software: you can Use, Study, Modify and Redistribute // This program is free software: you can Use, Study, Modify and Redistribute
...@@ -27,6 +27,7 @@ import ( ...@@ -27,6 +27,7 @@ import (
pickle "github.com/kisielk/og-rek" pickle "github.com/kisielk/og-rek"
"lab.nexedi.com/kirr/neo/go/zodb/internal/pickletools" "lab.nexedi.com/kirr/neo/go/zodb/internal/pickletools"
"lab.nexedi.com/kirr/neo/go/internal/zodbpickle"
"lab.nexedi.com/kirr/go123/xerr" "lab.nexedi.com/kirr/go123/xerr"
) )
...@@ -47,7 +48,7 @@ type PyData []byte ...@@ -47,7 +48,7 @@ type PyData []byte
// If pickle decoding fails - "?.?" is returned. // If pickle decoding fails - "?.?" is returned.
func (d PyData) ClassName() string { func (d PyData) ClassName() string {
// see ObjectReader.getClassName & get_pickle_metadata in zodb/py // see ObjectReader.getClassName & get_pickle_metadata in zodb/py
p := pickle.NewDecoder(bytes.NewReader([]byte(d))) p := zodbpickle.NewUnpickler(bytes.NewReader([]byte(d)), nil)
xklass, err := p.Decode() xklass, err := p.Decode()
if err != nil { if err != nil {
return "?.?" return "?.?"
...@@ -66,12 +67,7 @@ func encodePyData(pyclass pickle.Class, pystate interface{}) PyData { ...@@ -66,12 +67,7 @@ func encodePyData(pyclass pickle.Class, pystate interface{}) PyData {
// XXX better return mem.Buf instead of PyData? // XXX better return mem.Buf instead of PyData?
buf := &bytes.Buffer{} buf := &bytes.Buffer{}
p := pickle.NewEncoderWithConfig(buf, &pickle.EncoderConfig{ p := zodbpickle.NewPickler(buf, persistentRef)
// allow pristine python2 to decode the pickle.
// TODO 2 -> 3 since ZODB5 switched to it and uses zodbpickle.
Protocol: 2,
PersistentRef: persistentRef,
})
// emit: object type // emit: object type
err := p.Encode(pyclass) err := p.Encode(pyclass)
...@@ -99,10 +95,7 @@ func encodePyData(pyclass pickle.Class, pystate interface{}) PyData { ...@@ -99,10 +95,7 @@ func encodePyData(pyclass pickle.Class, pystate interface{}) PyData {
func (d PyData) decode(jar *Connection) (pyclass pickle.Class, pystate interface{}, err error) { func (d PyData) decode(jar *Connection) (pyclass pickle.Class, pystate interface{}, err error) {
defer xerr.Context(&err, "pydata: decode") defer xerr.Context(&err, "pydata: decode")
p := pickle.NewDecoderWithConfig( p := zodbpickle.NewUnpickler(bytes.NewReader([]byte(d)), jar.loadref)
bytes.NewReader([]byte(d)),
&pickle.DecoderConfig{PersistentLoad: jar.loadref},
)
xklass, err := p.Decode() xklass, err := p.Decode()
if err != nil { if err != nil {
......
...@@ -36,6 +36,7 @@ import ( ...@@ -36,6 +36,7 @@ import (
"lab.nexedi.com/kirr/neo/go/zodb/storage/fs1/fsb" "lab.nexedi.com/kirr/neo/go/zodb/storage/fs1/fsb"
pickle "github.com/kisielk/og-rek" pickle "github.com/kisielk/og-rek"
"lab.nexedi.com/kirr/neo/go/internal/zodbpickle"
"lab.nexedi.com/kirr/go123/mem" "lab.nexedi.com/kirr/go123/mem"
"lab.nexedi.com/kirr/go123/xbufio" "lab.nexedi.com/kirr/go123/xbufio"
...@@ -107,7 +108,7 @@ func (fsi *Index) Save(w io.Writer) (err error) { ...@@ -107,7 +108,7 @@ func (fsi *Index) Save(w io.Writer) (err error) {
err = &IndexSaveError{err} err = &IndexSaveError{err}
}() }()
p := pickle.NewEncoder(w) p := zodbpickle.NewPickler(w, nil)
err = p.Encode(fsi.TopPos) err = p.Encode(fsi.TopPos)
if err != nil { if err != nil {
...@@ -232,7 +233,7 @@ func LoadIndex(r io.Reader) (fsi *Index, err error) { ...@@ -232,7 +233,7 @@ func LoadIndex(r io.Reader) (fsi *Index, err error) {
xr := xbufio.NewReader(r) xr := xbufio.NewReader(r)
// by passing bufio.Reader directly we make sure it won't create one internally // by passing bufio.Reader directly we make sure it won't create one internally
p := pickle.NewDecoder(xr.Reader) p := zodbpickle.NewUnpickler(xr.Reader, nil)
picklePos = xr.InputOffset() picklePos = xr.InputOffset()
xtopPos, err = p.Decode() xtopPos, err = p.Decode()
......
...@@ -31,6 +31,7 @@ import ( ...@@ -31,6 +31,7 @@ import (
msgp "github.com/tinylib/msgp/msgp" msgp "github.com/tinylib/msgp/msgp"
msgpack "github.com/shamaton/msgpack" msgpack "github.com/shamaton/msgpack"
pickle "github.com/kisielk/og-rek" pickle "github.com/kisielk/og-rek"
"lab.nexedi.com/kirr/neo/go/internal/zodbpickle"
"lab.nexedi.com/kirr/go123/mem" "lab.nexedi.com/kirr/go123/mem"
"lab.nexedi.com/kirr/neo/go/zodb" "lab.nexedi.com/kirr/neo/go/zodb"
...@@ -84,7 +85,7 @@ func (e encoding) pktDecode(pkb *pktBuf) (msg, error) { ...@@ -84,7 +85,7 @@ func (e encoding) pktDecode(pkb *pktBuf) (msg, error) {
// pktEncodeZ encodes message into raw Z (pickle) packet. // pktEncodeZ encodes message into raw Z (pickle) packet.
func pktEncodeZ(m msg) *pktBuf { func pktEncodeZ(m msg) *pktBuf {
pkb := allocPkb() pkb := allocPkb()
p := pickle.NewEncoder(pkb) p := zodbpickle.NewPickler(pkb, nil)
err := p.Encode(pickle.Tuple{m.msgid, m.flags, m.method, m.arg}) err := p.Encode(pickle.Tuple{m.msgid, m.flags, m.method, m.arg})
if err != nil { if err != nil {
panic(err) // all our types are expected to be supported by pickle panic(err) // all our types are expected to be supported by pickle
...@@ -119,7 +120,7 @@ func pktEncodeM(m msg) *pktBuf { ...@@ -119,7 +120,7 @@ func pktEncodeM(m msg) *pktBuf {
func pktDecodeZ(pkb *pktBuf) (msg, error) { func pktDecodeZ(pkb *pktBuf) (msg, error) {
var m msg var m msg
// must be (msgid, False|0, ".reply", res) // must be (msgid, False|0, ".reply", res)
d := pickle.NewDecoder(bytes.NewReader(pkb.Payload())) d := zodbpickle.NewUnpickler(bytes.NewReader(pkb.Payload()), nil)
xpkt, err := d.Decode() xpkt, err := d.Decode()
if err != nil { if err != nil {
return m, err return m, err
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment