86 lines
3.2 KiB
Python
86 lines
3.2 KiB
Python
#
|
|
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Example HTTP operator and sensor"""
|
|
from __future__ import annotations
|
|
|
|
import gzip
|
|
import json
|
|
import os
|
|
from datetime import datetime
|
|
|
|
from airflow.models.baseoperator import chain
|
|
import requests
|
|
from airflow import DAG
|
|
from airflow.providers.http.operators.http import HttpOperator
|
|
from airflow.providers.http.sensors.http import HttpSensor
|
|
from airflow.decorators import task
|
|
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
|
|
|
|
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME", "zenodo-bucket")
|
|
S3_BUCKET_KEY = os.getenv("S3_BUCKET_KEY", "test")
|
|
S3_BUCKET_KEY_LIST = os.getenv("S3_BUCKET_KEY_LIST", "test2")
|
|
S3_BUCKET_WILDCARD_KEY = os.getenv("S3_BUCKET_WILDCARD_KEY", "test*")
|
|
PREFIX = os.getenv("S3_PREFIX", "test")
|
|
INACTIVITY_PERIOD = float(os.getenv("INACTIVITY_PERIOD", 5))
|
|
AWS_DEFAULT_REGION = os.getenv("AWS_DEFAULT_REGION", "us-east-1")
|
|
LOCAL_FILE_PATH = os.getenv("LOCAL_FILE_PATH", "/usr/local/airflow/dags/example_s3_test_file.txt")
|
|
AWS_CONN_ID = os.getenv("ASTRO_AWS_S3_CONN_ID", "s3_conn")
|
|
EXECUTION_TIMEOUT = int(os.getenv("EXECUTION_TIMEOUT", 6))
|
|
|
|
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
|
|
DAG_ID = "populate_opensearch"
|
|
|
|
|
|
@task
|
|
def bulk_load():
|
|
hook = S3Hook(AWS_CONN_ID, transfer_config_args={'use_threads': False})
|
|
|
|
keys = hook.list_keys(bucket_name=S3_BUCKET_NAME, prefix='/organization/')
|
|
session = requests.Session()
|
|
session.auth = ("admin", "admin")
|
|
|
|
session.put("https://opensearch-cluster.lot1-opensearch-cluster.svc.local/organization_index",
|
|
data=json.dumps({
|
|
"settings": {
|
|
"number_of_shards": 1,
|
|
"number_of_replicas": 1
|
|
}
|
|
}),
|
|
headers={"Content-Type": "application/json"},
|
|
verify=False)
|
|
|
|
for key in keys:
|
|
s3_obj = hook.get_key(key, bucket_name=S3_BUCKET_NAME)
|
|
with gzip.GzipFile(fileobj=s3_obj.get()["Body"]) as gzipfile:
|
|
for line in gzipfile:
|
|
session.post("https://opensearch-cluster.lot1-opensearch-cluster.svc.local/organization_index/_doc",
|
|
data=line,
|
|
headers={"Content-Type": "application/json"},
|
|
verify=False)
|
|
|
|
with DAG(
|
|
DAG_ID,
|
|
default_args={"retries": 1},
|
|
tags=["example"],
|
|
start_date=datetime(2021, 1, 1),
|
|
catchup=False,
|
|
) as dag:
|
|
bulk_load_data = bulk_load()
|
|
chain(bulk_load_data)
|
|
|