Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

flatten all level #6844

Merged
merged 10 commits into from
Apr 25, 2024
Merged

flatten all level #6844

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
db8e00f
flatten all level
Apr 2, 2024
1593d7a
make lint happy
Apr 3, 2024
0992db7
add option to enable/disable flatten
Apr 7, 2024
48d7782
trun to str
Apr 7, 2024
13a2760
debug line
Apr 7, 2024
0b86bc3
test case init
Apr 7, 2024
d2077b1
add test caces
Apr 7, 2024
4522448
Merge branch 'master' into kim/mongo_flatten_dict
KimBioInfoStudio Apr 7, 2024
8a78005
make lint happy
Apr 7, 2024
f15515d
Merge branch 'master' into kim/mongo_flatten_dict
justinclift Apr 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 45 additions & 14 deletions redash/query_runner/mongodb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ def datetime_parser(dct):
return bson_object_hook(dct, json_options=opts)


def parse_query_json(query):
def parse_query_json(query: str):
query_data = json_loads(query, object_hook=datetime_parser)
return query_data

Expand All @@ -82,26 +82,40 @@ def _get_column_by_name(columns, column_name):
return None


def _parse_dict(dic):
def _parse_dict(dic: dict, flatten: bool = False) -> dict:
res = {}
for key, value in dic.items():
if isinstance(value, dict):
for tmp_key, tmp_value in _parse_dict(value).items():
new_key = "{}.{}".format(key, tmp_key)
res[new_key] = tmp_value

def _flatten(x, name=""):
if isinstance(x, dict):
for k, v in x.items():
_flatten(v, "{}.{}".format(name, k))
elif isinstance(x, list):
for idx, item in enumerate(x):
_flatten(item, "{}.{}".format(name, idx))
else:
res[key] = value
res[name[1:]] = x

if flatten:
_flatten(dic)
else:
for key, value in dic.items():
if isinstance(value, dict):
for tmp_key, tmp_value in _parse_dict(value).items():
new_key = "{}.{}".format(key, tmp_key)
res[new_key] = tmp_value
else:
res[key] = value
return res


def parse_results(results):
def parse_results(results: list, flatten: bool = False) -> list:
rows = []
columns = []

for row in results:
parsed_row = {}

parsed_row = _parse_dict(row)
parsed_row = _parse_dict(row, flatten)
for column_name, value in parsed_row.items():
columns.append(
{
Expand Down Expand Up @@ -140,6 +154,14 @@ def configuration_schema(cls):
],
"title": "Replica Set Read Preference",
},
"flatten": {
"type": "string",
"extendedEnum": [
{"value": "False", "name": "False"},
{"value": "True", "name": "True"},
],
"title": "Flatten Results",
},
},
"secret": ["password"],
"required": ["connectionString", "dbName"],
Expand All @@ -160,6 +182,9 @@ def __init__(self, configuration):
True if "replicaSetName" in self.configuration and self.configuration["replicaSetName"] else False
)

self.flatten = self.configuration.get("flatten", "False").upper() in ["TRUE", "YES", "ON", "1", "Y", "T"]
logger.debug("flatten: {}".format(self.flatten))

@classmethod
def custom_json_encoder(cls, dec, o):
if isinstance(o, ObjectId):
Expand Down Expand Up @@ -278,8 +303,10 @@ def run_query(self, query, user): # noqa: C901
if "$sort" in step:
sort_list = []
for sort_item in step["$sort"]:
sort_list.append((sort_item["name"], sort_item["direction"]))

if isinstance(sort_item, dict):
sort_list.append((sort_item["name"], sort_item.get("direction", 1)))
elif isinstance(sort_item, list):
sort_list.append(tuple(sort_item))
step["$sort"] = SON(sort_list)

if "fields" in query_data:
Expand All @@ -289,7 +316,10 @@ def run_query(self, query, user): # noqa: C901
if "sort" in query_data and query_data["sort"]:
s = []
for field_data in query_data["sort"]:
s.append((field_data["name"], field_data["direction"]))
if isinstance(field_data, dict):
s.append((field_data["name"], field_data.get("direction", 1)))
elif isinstance(field_data, list):
s.append(tuple(field_data))

columns = []
rows = []
Expand Down Expand Up @@ -330,7 +360,7 @@ def run_query(self, query, user): # noqa: C901

rows.append({"count": cursor})
else:
rows, columns = parse_results(cursor)
rows, columns = parse_results(cursor, flatten=self.flatten)

if f:
ordered_columns = []
Expand All @@ -340,6 +370,7 @@ def run_query(self, query, user): # noqa: C901
ordered_columns.append(column)

columns = ordered_columns
logger.debug("columns: {}".format(columns))

if query_data.get("sortColumns"):
reverse = query_data["sortColumns"] == "desc"
Expand Down
56 changes: 55 additions & 1 deletion tests/query_runner/test_mongodb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,13 @@ def test_parses_nested_results(self):
"column": 2,
"column2": "test",
"column3": "hello",
"nested": {"a": 2, "b": "str2", "c": "c", "d": {"e": 3}},
"nested": {
"a": 2,
"b": "str2",
"c": "c",
"d": {"e": 3},
"f": {"h": {"i": ["j", "k", "l"]}},
},
},
]

Expand All @@ -158,6 +164,7 @@ def test_parses_nested_results(self):
"nested.b": "str2",
"nested.c": "c",
"nested.d.e": 3,
"nested.f.h.i": ["j", "k", "l"],
},
)

Expand All @@ -167,3 +174,50 @@ def test_parses_nested_results(self):
self.assertIsNotNone(_get_column_by_name(columns, "nested.a"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.b"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.c"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.d.e"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.f.h.i"))

def test_parses_flatten_nested_results(self):
raw_results = [
{
"column": 2,
"column2": "test",
"column3": "hello",
"nested": {
"a": 2,
"b": "str2",
"c": "c",
"d": {"e": 3},
"f": {"h": {"i": ["j", "k", "l"]}},
},
}
]

rows, columns = parse_results(raw_results, flatten=True)
print(rows)
self.assertDictEqual(
rows[0],
{
"column": 2,
"column2": "test",
"column3": "hello",
"nested.a": 2,
"nested.b": "str2",
"nested.c": "c",
"nested.d.e": 3,
"nested.f.h.i.0": "j",
"nested.f.h.i.1": "k",
"nested.f.h.i.2": "l",
},
)

self.assertIsNotNone(_get_column_by_name(columns, "column"))
self.assertIsNotNone(_get_column_by_name(columns, "column2"))
self.assertIsNotNone(_get_column_by_name(columns, "column3"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.a"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.b"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.c"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.d.e"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.f.h.i.0"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.f.h.i.1"))
self.assertIsNotNone(_get_column_by_name(columns, "nested.f.h.i.2"))
Loading