User Job -- Manual

gistfile1.txt

export COSMOS_GRAPH_KEY=your_cosmos_graph_key
export COSMOS_GRAPH_NAME=your_cosmos_graph_name
export DATA_PARTITION=your_data_partition
export USER_ID=your_user_id

kubectl apply -f - <<EOF
apiVersion: v1
kind: ConfigMap
metadata:
  name: add-user-script
  namespace: dcp
data:
  add-user.py: |
    #!/usr/bin/env python3
    from gremlin_python.driver import client, serializer, protocol
    from gremlin_python.driver.protocol import GremlinServerError
    import sys
    import asyncio
    import os

    def main():
        cosmosGrahName = os.getenv('COSMOS_GRAPH_NAME')
        cosmosGraphKey = os.getenv('COSMOS_GRAPH_KEY')
        dataPartition = os.getenv('DATA_PARTITION')
        userId = os.getenv('USER')

        print(f"Cosmos: {cosmosGrahName}")
        print(f"Data Partition: {dataPartition}")
        print(f"User: {userId}")

        if not cosmosGrahName or not cosmosGraphKey or not dataPartition or not userId:
            print("Missing required environment variables")
            sys.exit(1)

        if sys.platform == "win32":
            asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
        
        cosmosClient = client.Client(
            url=f"wss://{cosmosGrahName}.gremlin.cosmos.azure.com:443/", 
            traversal_source='g',
            username="/dbs/osdu-graph/colls/Entitlements",
            password=cosmosGraphKey,
            message_serializer=serializer.GraphSONSerializersV2d0()
        )

        if shouldCreateUser(cosmosClient, dataPartition, userId):
            print("Creating user")
            createUser(cosmosClient, dataPartition, userId)
        else:
            print("User already exists")
        
        cosmosClient.close()
            
    def shouldCreateUser(client, dataPartition, userId):
        result = client.submit(
            (
                f"g.V().has('nodeId', '{userId}')"
                f".has('dataPartitionId', '{dataPartition}')"
                f".out('parent')"
                f".has('name', 'users.datalake.ops')"
                f".has('dataPartitionId', '{dataPartition}')"
            )
        )    
        return not result.all().result()

    def createUser(client, dataPartition, userId):
        client.submit(
            (
                f"g.addV('USER').property('nodeId', '{userId}').property('dataPartitionId', '{dataPartition}').as('user')"
                f".V().has('dataPartitionId', '{dataPartition}').has('nodeId', 'users.datalake.ops@{dataPartition}.dataservices.energy').as('opsgroup')"
                f".addE('parent').from('user').to('opsgroup')"
                f".addE('child').from('opsgroup').to('user')"
                f".V().has('dataPartitionId', '{dataPartition}').has('nodeId', 'users@{dataPartition}.dataservices.energy').as('usersgroup')"
                f".addE('parent').from('user').to('usersgroup')"
                f".addE('child').from('usersgroup').to('user')"
            )
        )

        print("Created user")

    main()
---
apiVersion: batch/v1
kind: Job
metadata:
  name: add-user-job
  namespace: dcp
spec:
  template:
    spec:
      containers:
      - name: add-user
        image: mcr.microsoft.com/cbl-mariner/base/python:3.9
        env:
        - name: COSMOS_GRAPH_KEY
          value: ${COSMOS_GRAPH_KEY}
        - name: COSMOS_GRAPH_NAME
          value: ${COSMOS_GRAPH_NAME}
        - name: DATA_PARTITION
          value: ${DATA_PARTITION}
        - name: USER
          value: ${USER_ID}
        volumeMounts:
        - name: script-volume
          mountPath: /scripts
        command: ["sh", "-c", "pip install gremlinpython==3.7.2 && python3 /scripts/add-user.py"]
      restartPolicy: Never
      volumes:
      - name: script-volume
        configMap:
          name: add-user-script
EOF