Source code for ax.storage.sqa_store.json

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

import json
from typing import Any, Dict, List, Optional

from ax.storage.sqa_store.db import JSON_FIELD_LENGTH, MEDIUMTEXT_BYTES
from sqlalchemy.ext.mutable import MutableDict, MutableList
from sqlalchemy.types import VARCHAR, Text, TypeDecorator


[docs]class JSONEncodedObject(TypeDecorator):
    """Class for JSON-encoding objects in SQLAlchemy.

    Represents an object that is automatically marshalled and unmarshalled
    to/from the corresponding JSON string. By itself, this data structure does
    not track any changes.

    """

    impl: VARCHAR = VARCHAR(JSON_FIELD_LENGTH)

    def __init__(
        self, object_pairs_hook: Any = None, *args: List[Any], **kwargs: Dict[Any, Any]
    ) -> None:
        self.object_pairs_hook: Any = object_pairs_hook
        super().__init__(*args, **kwargs)

[docs]    def process_bind_param(self, value: Any, dialect: Any) -> Optional[str]:
        if value is not None:
            return json.dumps(value)
        else:
            return None

[docs]    def process_result_value(self, value: Any, dialect: Any) -> Any:
        if value is not None:
            return json.loads(value, object_pairs_hook=self.object_pairs_hook)
        else:
            return None


[docs]class JSONEncodedText(JSONEncodedObject):
    """Class for JSON-encoding objects in SQLAlchemy, backed by TEXT (MySQL).

    See description in JSONEncodedObject.

    """

    impl = Text


[docs]class JSONEncodedMediumText(JSONEncodedObject):
    """Class for JSON-encoding objects in SQLAlchemy, backed by MEDIUMTEXT
    (MySQL).

    See description in JSONEncodedObject.

    """

    impl = Text(MEDIUMTEXT_BYTES)


JSONEncodedList = MutableList.as_mutable(JSONEncodedObject)
JSONEncodedDict = MutableDict.as_mutable(JSONEncodedObject)
JSONEncodedTextDict = MutableDict.as_mutable(JSONEncodedText)