1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18from __future__ import annotations
19
20from typing import TYPE_CHECKING, Sequence
21
22from airflow.models import BaseOperator
23from airflow.providers.microsoft.azure.hooks.cosmos import AzureCosmosDBHook
24
25if TYPE_CHECKING:
26 from airflow.utils.context import Context
27
28
29class AzureCosmosInsertDocumentOperator(BaseOperator):
30 """Insert a new document into the specified Cosmos database and collection.
31
32 Both the database and collection will be created automatically if they do
33 not already exist.
34
35 :param database_name: The name of the database. (templated)
36 :param collection_name: The name of the collection. (templated)
37 :param document: The document to insert
38 :param azure_cosmos_conn_id: Reference to the
39 :ref:`Azure CosmosDB connection<howto/connection:azure_cosmos>`.
40 """
41
42 template_fields: Sequence[str] = ("database_name", "collection_name")
43 ui_color = "#e4f0e8"
44
45 def __init__(
46 self,
47 *,
48 database_name: str,
49 collection_name: str,
50 document: dict,
51 azure_cosmos_conn_id: str = "azure_cosmos_default",
52 **kwargs,
53 ) -> None:
54 super().__init__(**kwargs)
55 self.database_name = database_name
56 self.collection_name = collection_name
57 self.document = document
58 self.azure_cosmos_conn_id = azure_cosmos_conn_id
59
60 def execute(self, context: Context) -> None:
61 # Create the hook
62 hook = AzureCosmosDBHook(azure_cosmos_conn_id=self.azure_cosmos_conn_id)
63
64 # Create the DB if it doesn't already exist
65 if not hook.does_database_exist(self.database_name):
66 hook.create_database(self.database_name)
67
68 # Create the collection as well
69 if not hook.does_collection_exist(self.collection_name, self.database_name):
70 hook.create_collection(self.collection_name, self.database_name)
71
72 # finally insert the document
73 hook.upsert_document(self.document, self.database_name, self.collection_name)