Add support for CosmosDB bindings

Author: elprans

.travis.yml

@@ -17,6 +17,7 @@ addons:
 
         packages:
             - dotnet-sdk-2.1.4
+            - azure-functions-core-tools
 
 cache:
     pip: true

appveyor.yml

@@ -7,6 +7,9 @@ environment:
     matrix:
         - PYTHON: "C:\\Python36-x64\\python.exe"
 
+install:
+    - choco install azure-functions-core-tools --pre
+
 build_script:
     - "%PYTHON% -m pip install -U -e .[dev]"
     - "%PYTHON% setup.py webhost"

azure/functions/__init__.py

@@ -1,3 +1,4 @@
 from ._abc import HttpRequest, TimerRequest, InputStream, Context, Out  # NoQA
+from ._cosmosdb import Document, DocumentList  # NoQA
 from ._http import HttpResponse  # NoQA
 from ._queue import QueueMessage  # NoQA

azure/functions/_abc.py

@@ -197,3 +197,31 @@ def time_next_visible(self) -> typing.Optional[datetime.datetime]:
     @abc.abstractmethod
     def pop_receipt(self) -> typing.Optional[str]:
         pass
+
+
+class Document(abc.ABC):
+    @classmethod
+    @abc.abstractmethod
+    def from_json(cls, json_data: str) -> 'Document':
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def from_dict(cls, dct: dict) -> 'Document':
+        pass
+
+    @abc.abstractmethod
+    def __getitem__(self, key):
+        pass
+
+    @abc.abstractmethod
+    def __setitem__(self, key, value):
+        pass
+
+    @abc.abstractmethod
+    def to_json(self) -> str:
+        pass
+
+
+class DocumentList(abc.ABC):
+    pass

azure/functions/_cosmosdb.py

@@ -0,0 +1,46 @@
+import collections
+import json
+
+from . import _abc
+
+
+# Internal properties of CosmosDB documents.
+_SYSTEM_KEYS = {'_etag', '_lsn', '_rid', '_self', '_ts'}
+
+
+class Document(_abc.Document, collections.UserDict):
+    """An Azure Document.
+
+    Document objects are ``UserDict`` subclasses and behave like dicts.
+    """
+
+    @classmethod
+    def from_json(cls, json_data: str) -> 'Document':
+        """Create a Document from a JSON string."""
+        return cls.from_dict(json.loads(json_data))
+
+    @classmethod
+    def from_dict(cls, dct: dict) -> 'Document':
+        """Create a Document from a dict object."""
+        filtered = {k: v for k, v in dct.items() if k not in _SYSTEM_KEYS}
+        return cls(filtered)  # type: ignore
+
+    def to_json(self) -> str:
+        """Return the JSON representation of the document."""
+        return json.dumps(dict(self))
+
+    def __getitem__(self, key):
+        return collections.UserDict.__getitem__(self, key)
+
+    def __setitem__(self, key, value):
+        return collections.UserDict.__setitem__(self, key, value)
+
+    def __repr__(self) -> str:
+        return (
+            f'<azure.Document at 0x{id(self):0x}>'
+        )
+
+
+class DocumentList(_abc.DocumentList, collections.UserList):
+    "A ``UserList`` subclass containing a list of :class:`~Document` objects"
+    pass

azure/worker/bindings/__init__.py

@@ -8,6 +8,7 @@
 # Import type implementations and converters
 # to get them registered and available:
 from . import blob  # NoQA
+from . import cosmosdb  # NoQA
 from . import http  # NoQA
 from . import queue  # NoQA
 from . import timer  # NoQA

azure/worker/bindings/cosmosdb.py

@@ -0,0 +1,76 @@
+import collections.abc
+import json
+import typing
+
+from azure.functions import _cosmosdb as cdb
+
+from . import meta
+from .. import protos
+
+
+class CosmosDBConverter(meta.InConverter, meta.OutConverter,
+                        binding='cosmosDB'):
+
+    @classmethod
+    def check_input_type_annotation(cls, pytype: type) -> bool:
+        return issubclass(pytype, cdb.DocumentList)
+
+    @classmethod
+    def check_output_type_annotation(cls, pytype: type) -> bool:
+        return issubclass(pytype, (cdb.DocumentList, cdb.Document))
+
+    @classmethod
+    def from_proto(cls, data: protos.TypedData, *,
+                   pytype: typing.Optional[type],
+                   trigger_metadata) -> cdb.DocumentList:
+        data_type = data.WhichOneof('data')
+
+        if data_type == 'string':
+            body = data.string
+
+        elif data_type == 'bytes':
+            body = data.bytes.decode('utf-8')
+
+        elif data_type == 'json':
+            body = data.json
+
+        else:
+            raise NotImplementedError(
+                f'unsupported queue payload type: {data_type}')
+
+        documents = json.loads(body)
+        if not isinstance(documents, list):
+            documents = [documents]
+
+        return cdb.DocumentList(
+            cdb.Document.from_dict(doc) for doc in documents)
+
+    @classmethod
+    def to_proto(cls, obj: typing.Any, *,
+                 pytype: typing.Optional[type]) -> protos.TypedData:
+        if isinstance(obj, cdb.Document):
+            data = cdb.DocumentList([obj])
+
+        elif isinstance(obj, cdb.DocumentList):
+            data = obj
+
+        elif isinstance(obj, collections.abc.Iterable):
+            data = cdb.DocumentList()
+
+            for doc in obj:
+                if not isinstance(doc, cdb.Document):
+                    raise NotImplementedError
+                else:
+                    data.append(doc)
+
+        else:
+            raise NotImplementedError
+
+        return protos.TypedData(
+            json=json.dumps([dict(d) for d in data])
+        )
+
+
+class CosmosDBTriggerConverter(CosmosDBConverter,
+                               binding='cosmosDBTrigger', trigger=True):
+    pass

azure/worker/bindings/meta.py

@@ -1,4 +1,5 @@
 import abc
+import datetime
 import enum
 import json
 import typing
@@ -131,6 +132,22 @@ def _decode_trigger_metadata_field(
                 data, python_type=python_type,
                 context=f'field {field!r} in trigger metadata')
 
+    @classmethod
+    def _parse_datetime_metadata(
+            cls, trigger_metadata: typing.Mapping[str, protos.TypedData],
+            field: str) -> typing.Optional[datetime.datetime]:
+
+        datetime_str = cls._decode_trigger_metadata_field(
+            trigger_metadata, field, python_type=str)
+
+        if datetime_str is None:
+            return None
+        else:
+            # UTC ISO 8601 assumed
+            dt = datetime.datetime.strptime(
+                datetime_str, '%Y-%m-%dT%H:%M:%S+00:00')
+            return dt.replace(tzinfo=datetime.timezone.utc)
+
 
 class InConverter(_BaseConverter, binding=None):
 

azure/worker/bindings/queue.py

@@ -94,22 +94,6 @@ def from_proto(cls, data: protos.TypedData, *,
                 trigger_metadata, 'PopReceipt', python_type=str)
         )
 
-    @classmethod
-    def _parse_datetime_metadata(
-            cls, trigger_metadata: typing.Mapping[str, protos.TypedData],
-            field: str) -> typing.Optional[datetime.datetime]:
-
-        datetime_str = cls._decode_trigger_metadata_field(
-            trigger_metadata, field, python_type=str)
-
-        if datetime_str is None:
-            return None
-        else:
-            # UTC ISO 8601 assumed
-            dt = datetime.datetime.strptime(
-                datetime_str, '%Y-%m-%dT%H:%M:%S+00:00')
-            return dt.replace(tzinfo=datetime.timezone.utc)
-
 
 class QueueMessageOutConverter(meta.OutConverter, binding='queue'):
 

azure/worker/testutils.py

@@ -36,6 +36,7 @@
 TESTS_ROOT = PROJECT_ROOT / 'tests'
 DEFAULT_WEBHOST_DLL_PATH = PROJECT_ROOT / 'build' / 'webhost' / \
     'Microsoft.Azure.WebJobs.Script.WebHost.dll'
+EXTENSIONS_PATH = PROJECT_ROOT / 'build' / 'extensions' / 'bin'
 FUNCS_PATH = TESTS_ROOT / 'http_functions'
 WORKER_PATH = pathlib.Path(__file__).parent.parent.parent
 WORKER_CONFIG = WORKER_PATH / '.testconfig'
@@ -94,7 +95,7 @@ def get_script_dir(cls):
 
     @classmethod
     def setUpClass(cls):
-        script_dir = cls.get_script_dir()
+        script_dir = pathlib.Path(cls.get_script_dir())
         if os.environ.get('PYAZURE_WEBHOST_DEBUG'):
             cls.host_stdout = None
         else:
@@ -103,6 +104,17 @@ def setUpClass(cls):
         cls.webhost = start_webhost(script_dir=script_dir,
                                     stdout=cls.host_stdout)
 
+        extensions = TESTS_ROOT / script_dir / 'bin'
+
+        if not extensions.exists():
+            if extensions.is_symlink():
+                extensions.unlink()
+
+            extensions.symlink_to(EXTENSIONS_PATH, target_is_directory=True)
+            cls.linked_extensions = True
+        else:
+            cls.linked_extensions = False
+
     @classmethod
     def tearDownClass(cls):
         cls.webhost.close()
@@ -112,6 +124,10 @@ def tearDownClass(cls):
             cls.host_stdout.close()
             cls.host_stdout = None
 
+        if cls.linked_extensions:
+            extensions = TESTS_ROOT / cls.get_script_dir() / 'bin'
+            extensions.unlink()
+
     def _run_test(self, test, *args, **kwargs):
         if self.host_stdout is None:
             test(self, *args, **kwargs)
@@ -475,6 +491,10 @@ def popen_webhost(*, stdout, stderr, script_root=FUNCS_PATH, port=None):
         if st:
             extra_env['AzureWebJobsStorage'] = st
 
+        cosmos = testconfig['azure'].get('cosmosdb_key')
+        if cosmos:
+            extra_env['AzureWebJobsCosmosDBConnectionString'] = cosmos
+
     if port is not None:
         extra_env['ASPNETCORE_URLS'] = f'http://*:{port}'