Merge pull request #391 from treasure-data/stp-364-queries-report

add download of queries report

scenarios/account_reporting/README.md

@@ -8,6 +8,7 @@ This scenario shows how you can ingest account data from Treasure Data API's for
 - The API base urls for your region can be found here: https://api-docs.treasuredata.com/en/overview/aboutendpoints/#treasure-data-api-baseurls
 - `target.database` - Sets the database the data will be ingested to. This database must exist prior to running the workflow.
 - `target.tables` - Sets the table names for each report.
+- `reports_to_run` - [true|false] Enable/disable the download of a report
 
 2. Upload the workflow with TD CLI.
 ```

scenarios/account_reporting/account_reporting.dig

@@ -5,12 +5,25 @@ _export:
   if>: ${reports_to_run.activations_list}
   _do:
     +get_activations:
-      py>: activations.get_list
+      py>: scripts.activations.get_list
       destination_db: ${target.database}
       destination_tbl: ${target.tables.activations_list}
       _env:
         TD_API_KEY: ${secret:td.apikey}
         TD_API_BASEURL: ${td_api_baseurl}
         CDP_API_BASEURL: ${cdp_api_baseurl}
+      docker:
+        image: "digdag/digdag-python:3.9"
+
++get_queries:
+  if>: ${reports_to_run.queries_list}
+  _do:
+    +get_activations:
+      py>: scripts.queries.get_list
+      destination_db: ${target.database}
+      destination_tbl: ${target.tables.queries_list}
+      _env:
+        TD_API_KEY: ${secret:td.apikey}
+        TD_API_BASEURL: ${td_api_baseurl}
       docker:
         image: "digdag/digdag-python:3.9"

scenarios/account_reporting/config.yml

@@ -5,6 +5,8 @@ target:
   database: account_reporting
   tables:
     activations_list: activations
+    queries_list: queries
 
 reports_to_run:
-  activations_list: true
+  activations_list: true
+  queries_list: true

scenarios/account_reporting/scripts:queries.py

@@ -0,0 +1,36 @@
+import os
+import pytd
+import requests
+import pandas as pd
+import json
+
+def get_list(destination_db, destination_tbl):
+  # get all queries
+  queries_list = get_queries_list()
+  print(f'{len(queries_list)} queries found')
+
+  # remove sql from result 
+  queries_df = pd.DataFrame(queries_list)
+  queries_df.drop(['query'],axis=1,inplace=True)
+
+  if queries_df.empty:
+    print('No queries found on account')
+  else:
+    # write queries to db
+    apikey = os.environ["TD_API_KEY"]
+    td_api_baseurl = os.environ["TD_API_BASEURL"]
+    client = pytd.Client(endpoint=td_api_baseurl,apikey=apikey,database=destination_db,default_engine='presto')
+    client.load_table_from_dataframe(queries_df,destination=destination_tbl,writer='bulk_import',if_exists='overwrite')
+
+def get_queries_list():
+    return td_get("/v3/schedule/list")
+
+def td_get(endpoint):
+  apikey = os.environ["TD_API_KEY"]
+  headers = {'Authorization': 'TD1 ' + apikey}
+
+  td_api_baseurl = os.environ["TD_API_BASEURL"]
+  request_url = td_api_baseurl + endpoint 
+
+  response = requests.get(url = request_url, headers = headers)
+  return (response.json())['schedules']