We Keep Coding

Convert python script to airflow dag

I have identified the below script as being really useful for anyone running Amazon Redshift:
#!/usr/bin/env python
from __future__ import print_function
'''
analyze-vacuum-schema.py
* Copyright 2015, Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file 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.
The Redshift Analyze Vacuum Utility gives you the ability to automate VACUUM and ANALYZE operations.
When run, it will analyze or vacuum an entire schema or individual tables. This Utility Analyzes
and Vacuums table(s) in a Redshift Database schema, based on certain parameters like unsorted,
stats off and size of the table and system alerts from stl_explain & stl_alert_event_log.
By turning on/off '--analyze-flag' and '--vacuum-flag' parameters, you can run it as 'vacuum-only'
or 'analyze-only' utility. This script can be scheduled to run VACUUM and ANALYZE as part of
regular maintenance/housekeeping activities, when there are less database activities (quiet period).
This script will:
1) Analyze a single table or tables in a schema based on,
a) Alerts from stl_explain & stl_alert_event_log.
b) 'stats off' metrics from SVV_TABLE_INFO.
2) Vacuum a single table or tables in a schema based on,
a) The alerts from stl_alert_event_log.
b) The 'unsorted' and 'size' metrics from SVV_TABLE_INFO.
c) Vacuum reindex to analyze the interleaved sort keys
Srinikri Amazon Web Services (2015)
11/21/2015 : Added support for vacuum reindex to analyze the interleaved sort keys.
09/01/2017 : Fixed issues with interleaved sort key tables per https://github.com/awslabs/amazon-redshift-utils/issues/184
11/09/2017 : Refactored to support running in AWS Lambda
14/12/2017 : Refactored to support a more sensible interface style with kwargs
'''
import os
import sys
import argparse
# add the lib directory to the sys path
try:
sys.path.append(os.path.join(os.path.dirname(__file__), "lib"))
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
except:
pass
import getopt
import analyze_vacuum
import config_constants
__version__ = ".9.2.1"
OK = 0
ERROR = 1
INVALID_ARGS = 2
NO_WORK = 3
TERMINATED_BY_USER = 4
NO_CONNECTION = 5
# setup cli args
parser = argparse.ArgumentParser()
parser.add_argument("--analyze-flag", dest="analyze_flag", default=True, type=bool,
help="Flag to turn ON/OFF ANALYZE functionality (True or False : Default = True ")
parser.add_argument("--max-unsorted-pct", dest="max_unsorted_pct",
help="Maximum unsorted percentage(% to consider a table for vacuum : Default = 50%")
parser.add_argument("--min-interleaved-cnt", dest="min_interleaved_cnt", type=int,
help="Minimum stv_interleaved_counts records to consider a table for vacuum reindex: Default = 0")
parser.add_argument("--min-interleaved-skew", dest="min_interleaved_skew",
help="Minimum index skew to consider a table for vacuum reindex: Default = 1.4")
parser.add_argument("--min-unsorted-pct", dest="min_unsorted_pct",
help="Minimum unsorted percentage(% to consider a table for vacuum : Default = 5%")
parser.add_argument("--stats-off-pct ", dest="stats_off_pct",
help="Minimum stats off percentage(% to consider a table for analyze : Default = 10%")
parser.add_argument("--table-name", dest="table_name",
help="A specific table to be Analyzed or Vacuumed if analyze-schema is not desired")
parser.add_argument("--vacuum-flag", dest="vacuum_flag", default=True, type=bool,
help="Flag to turn ON/OFF VACUUM functionality (True or False : Default = True")
parser.add_argument("--vacuum-parameter", dest="vacuum_parameter",
help="Vacuum parameters [ FULL | SORT ONLY | DELETE ONLY | REINDEX ] Default = FULL")
parser.add_argument("--blacklisted-tables", dest="blacklisted_tables", help="The tables we do not want to Vacuum")
parser.add_argument("--db-conn-opts", dest="db_conn_opts",
help="Additional connection options. name1=opt1[ name2=opt2]..")
parser.add_argument("--db-host", dest="db_host", required=True, help="The Cluster endpoint")
parser.add_argument("--db-port", dest="db_port", type=int, required=True,
help="The Cluster endpoint port : Default = 5439")
parser.add_argument("--db-pwd", dest="db_pwd", help="The Password for the Database User to connect to")
parser.add_argument("--db-user", dest="db_user", required=True, help="The Database User to connect to")
parser.add_argument("--debug ", dest="debug", default=False,
help="Generate Debug Output including SQL Statements being run")
parser.add_argument("--ignore-errors", dest="ignore_errors", default=True,
help="Ignore errors raised when running and continue processing")
parser.add_argument("--max-table-size-mb", dest="max_table_size_mb", type=int,
help="Maximum table size in MB : Default = 700*1024 MB")
parser.add_argument("--output-file", dest="output_file", help="The full path to the output file to be generated")
parser.add_argument("--predicate-cols", dest="predicate_cols", help="Analyze predicate columns only")
parser.add_argument("--query-group", dest="query_group", help="Set the query_group for all queries")
parser.add_argument("--require-ssl", dest="require_ssl", default=False,
help="Does the connection require SSL? (True | False")
parser.add_argument("--schema-name", dest="schema_name",
help="The Schema to be Analyzed or Vacuumed (REGEX: Default = public")
parser.add_argument("--slot-count", dest="slot_count", help="Modify the wlm_query_slot_count : Default = 1")
parser.add_argument("--suppress-cloudwatch", dest="suppress_cw",
help="Don't emit CloudWatch metrics for analyze or vacuum when set to True")
parser.add_argument("--db", dest="db", help="The Database to Use")
full_args = parser.parse_args()
parse_args = {}
# remove args that end up as None
for k, v in vars(full_args).items():
if v is not None:
parse_args[k] = v
def main(argv):
# get environmental args
args = {config_constants.DB_NAME: os.environ.get('PGDATABASE', None),
config_constants.DB_USER: os.environ.get('PGUSER', None),
config_constants.DB_PASSWORD: os.environ.get('PGPASSWORD', None),
config_constants.DB_HOST: os.environ.get('PGHOST', None),
config_constants.DB_PORT: os.environ.get('PGPORT', 5439)}
# add argparse args
args.update(parse_args)
if args.get(config_constants.OUTPUT_FILE) is not None:
sys.stdout = open(args.get(config_constants.OUTPUT_FILE), 'w')
# invoke the main method of the utility
result = analyze_vacuum.run_analyze_vacuum(**args)
if result is not None:
sys.exit(result)
else:
sys.exit(0)
if __name__ == "__main__":
main(sys.argv)
However, you can only "easily" schedule it from cron in an EC2 or similar scheduler. So, I have been trying to find a way to run it as an airflow dag.
I have found two similar questions on stackoverflow and I guess it is missing some import commands perhaps? like the below:
from airflow import DAG
from airflow.models import Variable
I'm hoping to find someone familiar enough with airflow to be able to help me find the required entries or point me in the right direction?
If so, I can create a fork from the original script:
https://github.com/awslabs/amazon-redshift-utils/blob/master/src/AnalyzeVacuumUtility/analyze-vacuum-schema.py
Which will assist others in the future with the same goal.

How about creating a new custom operator? It should accept all the cli arguments and then you can pass them to code from existing script. Here is some rough draft of what I would do:
from airflow.models import BaseOperator
class AnalyzeVacuumOperator(BaseOperator):
def __init__(
self,
*,
analyze_flag: bool = True,
max_unsorted_pct: float = None,
# ... here goes all other arguments from argparse
**kwargs
):
super().__init__(**kwargs)
self.args = {
config_constants.DB_NAME: os.environ.get('PGDATABASE', None),
# Here goes all the rest of args defined in original script
# ...
# Then we add new arguments configured by users
# which are now passed using cli (argparse).
# TODO: check what are current key names so we match with them
analyze_flag: analyze_flag,
max_unsorted_pct: max_unsorted_pct,
# ...
}
def execute(context):
args = self.args
# Here we copy paste from the existing script
if args.get(config_constants.OUTPUT_FILE) is not None:
sys.stdout = open(args.get(config_constants.OUTPUT_FILE), 'w')
# invoke the main method of the utility
result = analyze_vacuum.run_analyze_vacuum(**args)
if result is not None:
sys.exit(result)
else:
sys.exit(0)
Then you can use such operator in your DAGs:
with DAG(...) as dag:
AnalyzeVacuumOperator(
task_id="vacuum_task",
analyze_flag=True,
slot_count=3,
)
Mind that you may need to adjust imports and path changes as in original script.

How do I pass variables defined from argparse in config as a module to other scripts?

The main problem I have is accessing variables defined in a configuration file as a module in other scripts when those variables are defined by argparse inputs. If I encapsulate this code within a function, I will not have access to it in other scripts. Similarly, if I don't use a function, the configuration file will expect inputs whenever it is called as a module. I think global variables might be the answer, but I don't know how to set them up in this use case?
import shutil
import inspect
import configparser
import sys
import argparse
def cmd_inputs():
parser = argparse.ArgumentParser()
parser.add_argument('-dir','--geofloodhomedir',help="File path to directory that will hold the cfg file \
and the inputs and outputs folders. Default is the GeoNet3 directory.",
type=str)
parser.add_argument('-p','--project',
help="Folder within GIS, Hydraulics, and NWM directories \
of Inputs and Outputs. Default is 'my_project'.",
type=str)
parser.add_argument('-n','--DEM_name',
help="Name of Input DEM (without extension) and the prefix used for all \
project outputs. Default is 'dem'",
type=str)
parser.add_argument('--no_chunk',help="Dont batch process DEM's > 1.5 GB for Network Extraction. Default is to chunk DEMs larger than 1.5 GB.",
action='store_true')
global home_dir
args = parser.parse_args()
if args.geofloodhomedir:
home_dir = args.geofloodhomedir
if os.path.abspath(home_dir) == os.path.dirname(os.path.abspath(__file__)):
home_dir = os.getcwd()
print("hello")
print('GeoNet and GeoFlood Home Directory: ')
print(home_dir)
else:
ab_path = os.path.abspath(__file__)
home_dir = os.path.dirname(ab_path)
print('Using default GeoNet and GeoFlood home directory:')
print(home_dir)
if args.project:
project_name=args.project
print(f"Project Name: {project_name}")
else:
project_name='my_project'
print(f"Default Project Name: {project_name}")
if args.DEM_name:
dem_name = args.DEM_name
print(f'DEM Name: {dem_name}')
else:
dem_name='dem'
print(f'Default DEM: {dem_name}')
if not args.no_chunk:
print("Chunking DEMs > 1.5 GBs in Network Extraction script")
chunk = 1
else:
print("Not chunking input DEM or its products in Network Extraction script")
chunk = 0
config = configparser.ConfigParser()
config['Section']={'geofloodhomedir':home_dir,
'projectname':project_name,
'dem_name':dem_name,
'Chunk_DEM':chunk}
with open(os.path.join(home_dir,'GeoFlood.cfg'),'w') as configfile:
config.write(configfile)
if __name__=='__main__':
cmd_inputs()
print ("Configuration Complete")```

How to pass arguments to main function within Python module?

I have recently started learning more about Python Packages and Modules. I'm currently busy updating my existing modules so that that can be run as script or imported as a module into my other code. I'm not sure how to construct my input arguments within my module and pass them to the main() function within my module.
I've have written my my main() function and called it under if __name__ == '__main__' passing the input arguments. The inputs are currently hard coded to show what I'm trying to achieve. Any help in how to correctly construct my input arguments that the user will pass, which will then be passed onto the main function will be appreciated.
As mentioned I'm trying to be able to use the following as a script when used directly or imported as a module into my other code and run from there. If I import it as a module would I call the main() function when importing it? Is the following structure correct in how I have written the following? Any advice is appreciated.
'''
Created on March 12, 2017
Create a new ArcHydro Schema
File Geodatabase and Rasters
Folder
#author: PeterW
'''
# import site-packages and modules
import re
from pathlib import Path
import arcpy
# set environment settings
arcpy.env.overwriteOutput = True
def archydro_rasters_folder(workspace):
"""Create rasters folder directory
if it doens't already exist"""
model_name = Path(workspace).name
layers_name = re.sub(r"\D+", "Layers", model_name)
layers_folder = Path(workspace, layers_name)
if layers_folder.exists():
arcpy.AddMessage("Rasters folder: {0} exists".format(layers_name))
else:
layers_folder.mkdir(parents=True)
arcpy.AddMessage("Rasters folder {0} created".format(layers_name))
def archydro_fgdb_schema(workspace, schema, dem):
"""Create file geodatabase using XML
schema and set coordinate system based
on input DEM if it doesn't already exist"""
model_name = Path(workspace).name
fgdb = "{0}.gdb".format(model_name)
if arcpy.Exists(str(Path(workspace, fgdb))):
arcpy.AddMessage("{0} file geodatabase exists".format(fgdb))
else:
new_fgdb = arcpy.CreateFileGDB_management(str(workspace), fgdb)
import_type = "SCHEMA_ONLY"
config_keyword = "DEFAULTS"
arcpy.AddMessage("New {0} file geodatabase created".format(fgdb))
arcpy.ImportXMLWorkspaceDocument_management(new_fgdb, schema,
import_type,
config_keyword)
arcpy.AddMessage("ArcHydro schema imported")
projection = arcpy.Describe(dem).spatialReference
projection_name = projection.PCSName
feature_dataset = Path(workspace, fgdb, "Layers")
arcpy.DefineProjection_management(str(feature_dataset),
projection)
arcpy.AddMessage("Changed projection to {0}".format(projection_name))
def main(workspace, dem, schema):
"""main function to create rasters folder
and file geodatabase"""
archydro_rasters_folder(workspace)
archydro_fgdb_schema(schema, dem, workspace)
if __name__ == '__main__':
main(workspace = r"E:\Projects\2016\01_Bertrand_Small_Projects\G113268\ArcHydro\Model04",
dem = r"E:\Projects\2016\01_Bertrand_Small_Projects\G113268\ArcHydro\DEM2\raw",
schema = r"E:\Python\Masters\Schema\ESRI_UC12\ModelBuilder\Schema\Model01.xml")
Updated: 17/03/13
The following is my updated Python module based on Jonathan's suggestions:
'''
Created on March 12, 2017
Create a new ArcHydro Schema
File Geodatabase and Rasters
Folder
#author: PeterW
'''
# import site-packages and modules
import re
from pathlib import Path
import arcpy
import argparse
# set environment settings
arcpy.env.overwriteOutput = True
def rasters_directory(workspace):
"""Create rasters folder directory
if it doens't already exist"""
model_name = Path(workspace).name
layers_name = re.sub(r"\D+", "Layers", model_name)
layers_folder = Path(workspace, layers_name)
if layers_folder.exists():
arcpy.AddMessage("Rasters folder: {0} exists".format(layers_name))
else:
layers_folder.mkdir(parents=True)
arcpy.AddMessage("Rasters folder {0} created".format(layers_name))
def fgdb_schema(workspace, schema, dem):
"""Create file geodatabase using XML
schema and set coordinate system based
on input DEM if it doesn't already exist"""
model_name = Path(workspace).name
fgdb = "{0}.gdb".format(model_name)
if arcpy.Exists(str(Path(workspace, fgdb))):
arcpy.AddMessage("{0} file geodatabase exists".format(fgdb))
else:
new_fgdb = arcpy.CreateFileGDB_management(str(workspace), fgdb)
import_type = "SCHEMA_ONLY"
config_keyword = "DEFAULTS"
arcpy.AddMessage("New {0} file geodatabase created".format(fgdb))
arcpy.ImportXMLWorkspaceDocument_management(new_fgdb, schema,
import_type,
config_keyword)
arcpy.AddMessage("ArcHydro schema imported")
projection = arcpy.Describe(dem).spatialReference
projection_name = projection.PCSName
feature_dataset = Path(workspace, fgdb, "Layers")
arcpy.DefineProjection_management(str(feature_dataset),
projection)
arcpy.AddMessage("Changed projection to {0}".format(projection_name))
def model_schema(workspace, schema, dem):
"""Create model schema: rasters folder
and file geodatabase"""
rasters_directory(workspace)
fgdb_schema(schema, dem, workspace)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Create a ArcHydro schema')
parser.add_argument('--workspace', metavar='path', required=True,
help='the path to workspace')
parser.add_argument('--schema', metavar='path', required=True,
help='path to schema')
parser.add_argument('--dem', metavar='path', required=True,
help='path to dem')
args = parser.parse_args()
model_schema(workspace=args.workspace, schema=args.schema, dem=args.dem)

This looks correct to me, and yes if you're looking to use this as a module you would import main. Though, it would probably be better to name it in a more descriptive way.
To clarify how __main__ and the function main() works. When you execute a module it will have a name which is stored in __name__. If you execute the module stand alone as a script it will have the name __main__. If you execute it as part of a module ie import it into another module it will have the name of the module.
The function main() can be named anything you would like, and that wouldn't affect your program. It's commonly named main in small scripts but it's not a particularly good name if it's part of a larger body of code.
In terms letting a user to input arguments when running as a script I would look into either using argparse or click
An example of how argparse would work.
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='Create a ArcHydro schema')
parser.add_argument('--workspace', metavar='path', required=True,
help='the path to workspace')
parser.add_argument('--schema', metavar='path', required=True,
help='path to schema')
parser.add_argument('--dem', metavar='path', required=True,
help='path to dem')
args = parser.parse_args()
main(workspace=args.workspace, schema=args.schema, dem=args.dem)

Trying to make a dynamic host file for ansible in python

This is my first post here, so if there are any questions or if something is unlcear, don't hesitate to ask.
I am trying to use a dynamic host file so I can build multiple vagrant machines without having to manage the host file first. This is what I found online:
#!/usr/bin/env python
# Adapted from Mark Mandel's implementation
# https://github.com/ansible/ansible/blob/devel/plugins/inventory/vagrant.py
import argparse
import json
import paramiko
import subprocess
import sys
def parse_args():
parser = argparse.ArgumentParser(description="Vagrant inventory script")
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('--list', action='store_true')
group.add_argument('--host')
return parser.parse_args()
def list_running_hosts():
cmd = "vagrant status --machine-readable"
status = subprocess.check_output(cmd.split()).rstrip()
hosts = []
for line in status.split('\n'):
(_, host, key, value) = line.split(',')
if key == 'state' and value == 'running':
hosts.append(host)
return hosts
def get_host_details(host):
cmd = "vagrant ssh-config {}".format(host)
p = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE)
config = paramiko.SSHConfig()
config.parse(p.stdout)
c = config.lookup(host)
return {'ansible_ssh_host': c['hostname'],
'ansible_ssh_port': c['port'],
'ansible_ssh_user': c['user'],
'ansible_ssh_private_key_file': c['identityfile'][0]}
def main():
args = parse_args()
if args.list:
hosts = list_running_hosts()
json.dump({'vagrant': hosts}, sys.stdout)
else:
details = get_host_details(args.host)
json.dump(details, sys.stdout)
if __name__ == '__main__':
main()
However, when I run this I get the following error:
ERROR! The file inventory/vagrant.py is marked as executable, but failed to execute correctly. If this is not supposed to be an executable script, correct this with `chmod -x inventory/vagrant.py`.
ERROR! Inventory script (inventory/vagrant.py) had an execution error: Traceback (most recent call last):
File "/home/sebas/Desktop/playbooks/inventory/vagrant.py", line 52, in <module>
main()
File "/home/sebas/Desktop/playbooks/inventory/vagrant.py", line 45, in main
hosts = list_running_hosts()
File "/home/sebas/Desktop/playbooks/inventory/vagrant.py", line 24, in list_running_hosts
(_, host, key, value) = line.split(',')
ValueError: too many values to unpack
ERROR! inventory/vagrant.py:4: Expected key=value host variable assignment, got: argparse
does anybody know what I did wrong? Thank you guys in advance!

I guess the problem is that vagrant status command will work only inside a directory with a Vagrantfile, or if the ID of a target machine is specified.
To get the state of all active Vagrant environments on the system, vagrant global-status should be used instead. But global-status has a drawback: it uses a cache and does not actively verify the state of machines.
So to reliably determine the state, first we need to get the IDs of all VMs with vagrant global-status and then check these IDs with vagrant status ID.

Controlling VirtualBox from commandline with python

We are using python virtualbox API for controlling the virtualbox. For that we are using the "pyvb" package(as given in python API documentation).
al=pyvb.vb.VB()
m=pyvb.vm.vbVM()
al.startVM(m)
we have executed using the python interpreter. No error is shown but the virtualbox doesnt start. Could you please tell us what could be wrong(all necessary modules and packages have been imported)

I found that I can use the following functions to find if a VM is running, restore a VM to a specific snapshot, and start a VM by name.
from subprocess import Popen, PIPE
def running_vms():
"""
Return list of running vms
"""
f = Popen(r'vboxmanage --nologo list runningvms', stdout=PIPE).stdout
data = [ eachLine.strip() for eachLine in f ]
return data
def restore_vm(name='', snapshot=''):
"""
Restore VM to specific snapshot uuid
name = VM Name
snapshot = uuid of snapshot (uuid can be found in the xml file of your machines folder)
"""
command = r'vboxmanage --nologo snapshot %s restore %s' % (name,snapshot)
f = Popen(command, stdout=PIPE).stdout
data = [ eachLine.strip() for eachLine in f ]
return data
def launch_vm(name=''):
"""
Launch VM
name = VM Name
"""
command = r'vboxmanage --nologo startvm %s ' % name
f = Popen(command, stdout=PIPE).stdout
data = [ eachLine.strip() for eachLine in f ]
return data

The code quoted doesn't seem to specify what VM to run. Shouldn't you be doing a getVM call and then using that resulting VM instance in your startVM call? E.g.:
al=pyvb.vb.VB()
m=al.getVM(guid_of_vm)
al.startVM(m)
...would start the VM identified with the given GUID (all VirtualBox VMs have a GUID assigned when they're created). You can get the GUID from the VM's XML file. If you need to discover VMs at runtime, there's the handy listVMS call:
al=pyvb.vb.VB()
l=al.listVMS()
# choose a VM from the list, assign to 'm'
al.startVM(m)