How to create shard&index in airflow mongohook?

Viewed 20

I want to run mongo command with mongohook of airflow. How can I do it?

sh.shardCollection(db_name +, { _id: "hashed" }, false, { numInitialChunks: 128 });
db.collection.createIndex({ "field": 1 }, { field: true });
1 Answers

The pymongo client which the Mongohook provided in Airflow uses doesn't support the sh.shardCollection command in your script.

Though the createIndex collection method is supported in the pymongo client.

I recommend anyway to install the mongosh CLI binary and bake it into your container image for your workers.

You can write your shell command to a script such as /dags/templates/mongo-admin-create-index.js or some other location that it can be found.

Then can implement a custom operator using the SubprocessHook to run mongosh CLI command such as:

mongosh -f {mongosh_script} {db_address}

This custom operator would be along these lines

from airflow.compat.functools import cached_property
from airflow.hooks.subprocess import SubprocessHook
from airflow.providers.mongo.hooks import MongoHook

class MongoshScriptOperator(BaseOperator):

    template_fields: Sequence[str] = ('mongosh_script')

    def __init__(
        self,
        *,
        mongosh_script: str,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.mongosh_script = mongosh_script

    @cached_property
    def subprocess_hook(self):
        """Returns hook for running the shell command"""
        return SubprocessHook()

    def execute(self):
        """Executes a mongosh script"""
        mh = MongoHook(self.conn_id)
        
        self.subprocess_hook.run_command(
            command=['mongosh', '-f', self.mongosh_script, mh.uri],
        )

When creating the DagNode, you can pass the location of the script to your custom operator.

Related