Skip to content

Commit

Permalink
python: Replace ujson with orjson.
Browse files Browse the repository at this point in the history 
Fixes zulip#6507.

Signed-off-by: Anders Kaseorg <[email protected]>
  • Loading branch information
@andersk @timabbott
andersk authored and timabbott committed Aug 11, 2020
1 parent 123790a commit 61d0417
Show file tree
Hide file tree
Showing 132 changed files with 1,319 additions and 1,238 deletions.
4 changes: 2 additions & 2 deletions analytics/tests/test_counts.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from typing import Any, Dict, List, Optional, Tuple, Type
from unittest import mock

import ujson
import orjson
from django.apps import apps
from django.db import models
from django.db.models import Sum
Expand Down Expand Up @@ -1393,7 +1393,7 @@ def mark_audit_active(self, user: UserProfile, end_time: Optional[datetime]=None
end_time = self.TIME_ZERO
UserCount.objects.create(
user=user, realm=user.realm, property='active_users_audit:is_bot:day',
subgroup=ujson.dumps(user.is_bot), end_time=end_time, value=1)
subgroup=orjson.dumps(user.is_bot).decode(), end_time=end_time, value=1)

def mark_15day_active(self, user: UserProfile, end_time: Optional[datetime]=None) -> None:
if end_time is None:
Expand Down
4 changes: 2 additions & 2 deletions analytics/tests/test_views.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from typing import List, Optional
from unittest import mock

import ujson
import orjson
from django.http import HttpResponse
from django.utils.timezone import now as timezone_now

Expand Down Expand Up @@ -547,7 +547,7 @@ def check_realm_reactivation_link_query_result(result: HttpResponse) -> None:
stream_ids = [self.get_stream_id("Denmark")]
invitee_emails = [self.nonreg_email("test1")]
self.client_post("/json/invites", {"invitee_emails": invitee_emails,
"stream_ids": ujson.dumps(stream_ids),
"stream_ids": orjson.dumps(stream_ids).decode(),
"invite_as": PreregistrationUser.INVITE_AS['MEMBER']})
result = self.client_get("/activity/support", {"q": self.nonreg_email("test1")})
check_preregistration_user_query_result(result, self.nonreg_email("test1"), invite=True)
Expand Down
13 changes: 9 additions & 4 deletions corporate/lib/stripe.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
from functools import wraps
from typing import Callable, Dict, Optional, Tuple, TypeVar, cast

import orjson
import stripe
import ujson
from django.conf import settings
from django.core.signing import Signer
from django.db import transaction
Expand Down Expand Up @@ -297,10 +297,10 @@ def make_end_of_cycle_updates_if_needed(plan: CustomerPlan,
RealmAuditLog.objects.create(
realm=new_plan.customer.realm, event_time=event_time,
event_type=RealmAuditLog.CUSTOMER_SWITCHED_FROM_MONTHLY_TO_ANNUAL_PLAN,
extra_data=ujson.dumps({
extra_data=orjson.dumps({
"monthly_plan_id": plan.id,
"annual_plan_id": new_plan.id,
})
}).decode()
)
return new_plan, new_plan_ledger_entry

Expand Down Expand Up @@ -348,6 +348,11 @@ def compute_plan_parameters(
next_invoice_date = period_end
return billing_cycle_anchor, next_invoice_date, period_end, price_per_license

def decimal_to_float(obj: object) -> object:
if isinstance(obj, Decimal):
return float(obj)
raise TypeError # nocoverage

# Only used for cloud signups
@catch_stripe_errors
def process_initial_upgrade(user: UserProfile, licenses: int, automanage_licenses: bool,
Expand Down Expand Up @@ -424,7 +429,7 @@ def process_initial_upgrade(user: UserProfile, licenses: int, automanage_license
RealmAuditLog.objects.create(
realm=realm, acting_user=user, event_time=billing_cycle_anchor,
event_type=RealmAuditLog.CUSTOMER_PLAN_CREATED,
extra_data=ujson.dumps(plan_params))
extra_data=orjson.dumps(plan_params, default=decimal_to_float).decode())

if not free_trial:
stripe.InvoiceItem.create(
Expand Down
61 changes: 31 additions & 30 deletions corporate/tests/test_stripe.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, TypeVar, cast
from unittest.mock import Mock, patch

import orjson
import responses
import stripe
import ujson
from django.conf import settings
from django.core import signing
from django.http import HttpResponse
Expand Down Expand Up @@ -127,7 +127,8 @@ def read_stripe_fixture(decorated_function_name: str,
def _read_stripe_fixture(*args: Any, **kwargs: Any) -> Any:
mock = operator.attrgetter(mocked_function_name)(sys.modules[__name__])
fixture_path = stripe_fixture_path(decorated_function_name, mocked_function_name, mock.call_count)
fixture = ujson.load(open(fixture_path))
with open(fixture_path, "rb") as f:
fixture = orjson.loads(f.read())
# Check for StripeError fixtures
if "json_body" in fixture:
requestor = stripe.api_requestor.APIRequestor()
Expand Down Expand Up @@ -331,7 +332,7 @@ def upgrade(self, invoice: bool=False, talk_to_stripe: bool=True,
if key in params:
del params[key]
for key, value in params.items():
params[key] = ujson.dumps(value)
params[key] = orjson.dumps(value).decode()
return self.client_post("/json/billing/upgrade", params, **host_args)

# Upgrade without talking to Stripe
Expand Down Expand Up @@ -469,7 +470,7 @@ def test_upgrade_by_card(self, *mocks: Mock) -> None:
# TODO: Check for REALM_PLAN_TYPE_CHANGED
# (RealmAuditLog.REALM_PLAN_TYPE_CHANGED, Kandra()),
])
self.assertEqual(ujson.loads(RealmAuditLog.objects.filter(
self.assertEqual(orjson.loads(RealmAuditLog.objects.filter(
event_type=RealmAuditLog.CUSTOMER_PLAN_CREATED).values_list(
'extra_data', flat=True).first())['automanage_licenses'], True)
# Check that we correctly updated Realm
Expand Down Expand Up @@ -554,7 +555,7 @@ def test_upgrade_by_invoice(self, *mocks: Mock) -> None:
# TODO: Check for REALM_PLAN_TYPE_CHANGED
# (RealmAuditLog.REALM_PLAN_TYPE_CHANGED, Kandra()),
])
self.assertEqual(ujson.loads(RealmAuditLog.objects.filter(
self.assertEqual(orjson.loads(RealmAuditLog.objects.filter(
event_type=RealmAuditLog.CUSTOMER_PLAN_CREATED).values_list(
'extra_data', flat=True).first())['automanage_licenses'], False)
# Check that we correctly updated Realm
Expand Down Expand Up @@ -631,7 +632,7 @@ def test_free_trial_upgrade_by_card(self, *mocks: Mock) -> None:
# TODO: Check for REALM_PLAN_TYPE_CHANGED
# (RealmAuditLog.REALM_PLAN_TYPE_CHANGED, Kandra()),
])
self.assertEqual(ujson.loads(RealmAuditLog.objects.filter(
self.assertEqual(orjson.loads(RealmAuditLog.objects.filter(
event_type=RealmAuditLog.CUSTOMER_PLAN_CREATED).values_list(
'extra_data', flat=True).first())['automanage_licenses'], True)

Expand Down Expand Up @@ -785,7 +786,7 @@ def test_free_trial_upgrade_by_invoice(self, *mocks: Mock) -> None:
# TODO: Check for REALM_PLAN_TYPE_CHANGED
# (RealmAuditLog.REALM_PLAN_TYPE_CHANGED, Kandra()),
])
self.assertEqual(ujson.loads(RealmAuditLog.objects.filter(
self.assertEqual(orjson.loads(RealmAuditLog.objects.filter(
event_type=RealmAuditLog.CUSTOMER_PLAN_CREATED).values_list(
'extra_data', flat=True).first())['automanage_licenses'], False)

Expand Down Expand Up @@ -977,7 +978,7 @@ def test_upgrade_with_tampered_seat_count(self) -> None:
self.login_user(hamlet)
response = self.upgrade(talk_to_stripe=False, salt='badsalt')
self.assert_json_error_contains(response, "Something went wrong. Please contact")
self.assertEqual(ujson.loads(response.content)['error_description'], 'tampered seat count')
self.assertEqual(orjson.loads(response.content)['error_description'], 'tampered seat count')

def test_upgrade_race_condition(self) -> None:
hamlet = self.example_user('hamlet')
Expand All @@ -995,7 +996,7 @@ def check_error(error_description: str, upgrade_params: Mapping[str, Any],
del_args: Sequence[str] = []) -> None:
response = self.upgrade(talk_to_stripe=False, del_args=del_args, **upgrade_params)
self.assert_json_error_contains(response, "Something went wrong. Please contact")
self.assertEqual(ujson.loads(response.content)['error_description'], error_description)
self.assertEqual(orjson.loads(response.content)['error_description'], error_description)

hamlet = self.example_user('hamlet')
self.login_user(hamlet)
Expand All @@ -1015,13 +1016,13 @@ def check_min_licenses_error(invoice: bool, licenses: Optional[int], min_license
response = self.upgrade(invoice=invoice, talk_to_stripe=False,
del_args=del_args, **upgrade_params)
self.assert_json_error_contains(response, f"at least {min_licenses_in_response} users")
self.assertEqual(ujson.loads(response.content)['error_description'], 'not enough licenses')
self.assertEqual(orjson.loads(response.content)['error_description'], 'not enough licenses')

def check_max_licenses_error(licenses: int) -> None:
response = self.upgrade(invoice=True, talk_to_stripe=False,
licenses=licenses)
self.assert_json_error_contains(response, f"with more than {MAX_INVOICED_LICENSES} licenses")
self.assertEqual(ujson.loads(response.content)['error_description'], 'too many licenses')
self.assertEqual(orjson.loads(response.content)['error_description'], 'too many licenses')

def check_success(invoice: bool, licenses: Optional[int], upgrade_params: Dict[str, Any]={}) -> None:
if licenses is None:
Expand Down Expand Up @@ -1072,7 +1073,7 @@ def test_upgrade_with_uncaught_exception(self, mock_: Mock) -> None:
with patch("corporate.views.process_initial_upgrade", side_effect=Exception):
response = self.upgrade(talk_to_stripe=False)
self.assert_json_error_contains(response, "Something went wrong. Please contact [email protected].")
self.assertEqual(ujson.loads(response.content)['error_description'], 'uncaught exception during upgrade')
self.assertEqual(orjson.loads(response.content)['error_description'], 'uncaught exception during upgrade')

def test_request_sponsorship(self) -> None:
user = self.example_user("hamlet")
Expand All @@ -1081,9 +1082,9 @@ def test_request_sponsorship(self) -> None:
self.login_user(user)

data = {
"organization-type": ujson.dumps("Open-source"),
"website": ujson.dumps("https://infinispan.org/"),
"description": ujson.dumps("Infinispan is a distributed in-memory key/value data store with optional schema."),
"organization-type": orjson.dumps("Open-source").decode(),
"website": orjson.dumps("https://infinispan.org/").decode(),
"description": orjson.dumps("Infinispan is a distributed in-memory key/value data store with optional schema.").decode(),
}
response = self.client_post("/json/billing/sponsorship", data)
self.assert_json_success(response)
Expand Down Expand Up @@ -1281,10 +1282,10 @@ def test_replace_payment_source(self, *mocks: Mock) -> None:
with patch("corporate.lib.stripe.billing_logger.error") as mock_billing_logger:
with patch("stripe.Invoice.list") as mock_invoice_list:
response = self.client_post("/json/billing/sources/change",
{'stripe_token': ujson.dumps(stripe_token)})
{'stripe_token': orjson.dumps(stripe_token).decode()})
mock_billing_logger.assert_called()
mock_invoice_list.assert_not_called()
self.assertEqual(ujson.loads(response.content)['error_description'], 'card error')
self.assertEqual(orjson.loads(response.content)['error_description'], 'card error')
self.assert_json_error_contains(response, 'Your card was declined')
for stripe_source in stripe_get_customer(stripe_customer_id).sources:
assert isinstance(stripe_source, stripe.Card)
Expand All @@ -1295,9 +1296,9 @@ def test_replace_payment_source(self, *mocks: Mock) -> None:
stripe_token = stripe_create_token(card_number='4000000000000341').id
with patch("corporate.lib.stripe.billing_logger.error") as mock_billing_logger:
response = self.client_post("/json/billing/sources/change",
{'stripe_token': ujson.dumps(stripe_token)})
{'stripe_token': orjson.dumps(stripe_token).decode()})
mock_billing_logger.assert_called()
self.assertEqual(ujson.loads(response.content)['error_description'], 'card error')
self.assertEqual(orjson.loads(response.content)['error_description'], 'card error')
self.assert_json_error_contains(response, 'Your card was declined')
for stripe_source in stripe_get_customer(stripe_customer_id).sources:
assert isinstance(stripe_source, stripe.Card)
Expand All @@ -1309,7 +1310,7 @@ def test_replace_payment_source(self, *mocks: Mock) -> None:
# Replace with a valid card
stripe_token = stripe_create_token(card_number='5555555555554444').id
response = self.client_post("/json/billing/sources/change",
{'stripe_token': ujson.dumps(stripe_token)})
{'stripe_token': orjson.dumps(stripe_token).decode()})
self.assert_json_success(response)
number_of_sources = 0
for stripe_source in stripe_get_customer(stripe_customer_id).sources:
Expand Down Expand Up @@ -1435,8 +1436,8 @@ def test_switch_from_monthly_plan_to_annual_plan_for_automatic_license_managemen
self.assertEqual(annual_ledger_entries.values_list('licenses', 'licenses_at_next_renewal')[1], (25, 25))
audit_log = RealmAuditLog.objects.get(event_type=RealmAuditLog.CUSTOMER_SWITCHED_FROM_MONTHLY_TO_ANNUAL_PLAN)
self.assertEqual(audit_log.realm, user.realm)
self.assertEqual(ujson.loads(audit_log.extra_data)["monthly_plan_id"], monthly_plan.id)
self.assertEqual(ujson.loads(audit_log.extra_data)["annual_plan_id"], annual_plan.id)
self.assertEqual(orjson.loads(audit_log.extra_data)["monthly_plan_id"], monthly_plan.id)
self.assertEqual(orjson.loads(audit_log.extra_data)["annual_plan_id"], annual_plan.id)

invoice_plans_as_needed(self.next_month)

Expand Down Expand Up @@ -1839,15 +1840,15 @@ def test_who_can_access_json_endpoints(self) -> None:
self.login_user(self.example_user('hamlet'))
with patch("corporate.views.do_replace_payment_source") as mocked1:
response = self.client_post("/json/billing/sources/change",
{'stripe_token': ujson.dumps('token')})
{'stripe_token': orjson.dumps('token').decode()})
self.assert_json_success(response)
mocked1.assert_called()

# Realm owners have access, even if they are not billing admins
self.login_user(self.example_user('desdemona'))
with patch("corporate.views.do_replace_payment_source") as mocked2:
response = self.client_post("/json/billing/sources/change",
{'stripe_token': ujson.dumps('token')})
{'stripe_token': orjson.dumps('token').decode()})
self.assert_json_success(response)
mocked2.assert_called()

Expand All @@ -1858,21 +1859,21 @@ def verify_user_cant_access_endpoint(username: str, endpoint: str, request_data:
self.assert_json_error_contains(response, error_message)

verify_user_cant_access_endpoint("polonius", "/json/billing/upgrade",
{'billing_modality': ujson.dumps("charge_automatically"), 'schedule': ujson.dumps("annual"),
'signed_seat_count': ujson.dumps('signed count'), 'salt': ujson.dumps('salt')},
{'billing_modality': orjson.dumps("charge_automatically").decode(), 'schedule': orjson.dumps("annual").decode(),
'signed_seat_count': orjson.dumps('signed count').decode(), 'salt': orjson.dumps('salt').decode()},
"Must be an organization member")

verify_user_cant_access_endpoint("polonius", "/json/billing/sponsorship",
{'organization-type': ujson.dumps("event"), 'description': ujson.dumps("event description"),
'website': ujson.dumps("example.com")},
{'organization-type': orjson.dumps("event").decode(), 'description': orjson.dumps("event description").decode(),
'website': orjson.dumps("example.com").decode()},
"Must be an organization member")

for username in ["cordelia", "iago"]:
self.login_user(self.example_user(username))
verify_user_cant_access_endpoint(username, "/json/billing/sources/change", {'stripe_token': ujson.dumps('token')},
verify_user_cant_access_endpoint(username, "/json/billing/sources/change", {'stripe_token': orjson.dumps('token').decode()},
"Must be a billing administrator or an organization owner")

verify_user_cant_access_endpoint(username, "/json/billing/plan/change", {'status': ujson.dumps(1)},
verify_user_cant_access_endpoint(username, "/json/billing/plan/change", {'status': orjson.dumps(1).decode()},
"Must be a billing administrator or an organization owner")

# Make sure that we are testing all the JSON endpoints
Expand Down
4 changes: 1 addition & 3 deletions requirements/common.in
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,9 +101,7 @@ sourcemap
tornado==4.* # https://github.com/zulip/zulip/issues/8913

# Fast JSON parser
# Forked for an issue related to unpaired Unicode surrogates:
# https://github.com/zulip/zulip/issues/6332, https://github.com/esnme/ultrajson/pull/284
https://github.com/zulip/ultrajson/archive/70ac02becc3e11174cd5072650f885b30daab8a8.zip#egg=ujson==1.35+git
orjson

# Django extension for serving webpack modules
django-webpack4-loader
Expand Down
Loading

0 comments on commit 61d0417

Please sign in to comment.