Skip to content

Commit

Permalink
fix client pickle error when using password
Browse files Browse the repository at this point in the history
  • Loading branch information
wangyinz committed Jun 17, 2024
1 parent cc4a153 commit b744edf
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 2 deletions.
47 changes: 45 additions & 2 deletions python/mspasspy/db/client.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import pymongo
from typing import Any
from mspasspy.db.database import Database


Expand All @@ -16,9 +17,51 @@ class DBClient(pymongo.MongoClient):
creating a MongoClient or the MsPASS DBClient (this class).
"""

def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __init__(self, host=None, *args, **kwargs):
super().__init__(host=host, *args, **kwargs)
self.__default_database_name = self._MongoClient__default_database_name
self._mspass_db_host=host

def _repr_helper(self) -> str:
def option_repr(option: str, value: Any) -> str:
"""Fix options whose __repr__ isn't usable in a constructor."""
if option == "document_class":
if value is dict:
return "document_class=dict"
else:
return f"document_class={value.__module__}.{value.__name__}"
if option in pymongo.common.TIMEOUT_OPTIONS and value is not None:
return f"{option}={int(value * 1000)}"

return f"{option}={value!r}"

# Host first...
if self._mspass_db_host:
options = [
"host='{}'".format(self._mspass_db_host)
]
else:
options = [
"host=%r"
% [
"%s:%d" % (host, port) if port is not None else host
for host, port in self._topology_settings.seeds
]
]
# ... then everything in self._constructor_args...
options.extend(
option_repr(key, self.options._options[key]) for key in self._constructor_args
)
# ... then everything else.
options.extend(
option_repr(key, self.options._options[key])
for key in self.options._options
if key not in set(self._constructor_args) and key != "username" and key != "password"
)
return ", ".join(options)

def __repr__(self) -> str:
return f"{type(self).__name__}({self._repr_helper()})"

def __getitem__(self, name):
"""
Expand Down
1 change: 1 addition & 0 deletions python/mspasspy/db/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@ def __setstate__(self, data):
# work without it. Not sure how the symbol MongoClient is required
# here but it is - ignore if a lint like ide says MongoClient is not used
from pymongo import MongoClient
from mspasspy.db.client import DBClient

# The following is also needed for this object to be serialized correctly
# with dask distributed. Otherwise, the deserialized codec_options
Expand Down

0 comments on commit b744edf

Please sign in to comment.