| 1 |
#!/home/jsoc/bin/linux_x86_64/activepython |
| 2 |
|
| 3 |
import sys |
| 4 |
import os.path |
| 5 |
import getopt |
| 6 |
import pwd |
| 7 |
import re |
| 8 |
import psycopg2 |
| 9 |
|
| 10 |
# Return codes |
| 11 |
RET_SUCCESS = 0 |
| 12 |
RET_INVALIDARG = 1 |
| 13 |
RET_DBCONNECT = 2 |
| 14 |
RET_SQL = 3 |
| 15 |
|
| 16 |
def GetArgs(args): |
| 17 |
istat = bool(0) |
| 18 |
optD = {} |
| 19 |
|
| 20 |
try: |
| 21 |
opts, remainder = getopt.getopt(args, "hs:t:k:u:n:h:p:d", ["source=", "targets=", "keys=", "unique=", "dbname=", "dbhost=", "dbport="]) |
| 22 |
except getopt.GetoptError: |
| 23 |
print('Usage:\n fixsharp2.py [-h] -s <source series> -t <target-series list> -k <keys> -u <prime key key-list> -n <db name> -h <db host> -p <db port> [-d]', file=sys.stderr) |
| 24 |
istat = bool(1) |
| 25 |
|
| 26 |
if istat == bool(0): |
| 27 |
for opt, arg in opts: |
| 28 |
if opt == '-h': |
| 29 |
print('Usage:\n fixsharp2.py [-h] -s <source series> -t <target-series list> -k <keys> -u <prime key key-list> -n <db name> -h <db host> -p <db port> [-d]') |
| 30 |
sys.exit(0) |
| 31 |
elif opt in ("-s", "--source"): |
| 32 |
regexp = re.compile(r"\s*(\S+)\.(\S+)\s*") |
| 33 |
matchobj = regexp.match(arg) |
| 34 |
if matchobj is None: |
| 35 |
istat = bool(1) |
| 36 |
else: |
| 37 |
optD['ns'] = matchobj.group(1) |
| 38 |
optD['table'] = matchobj.group(2) |
| 39 |
elif opt in ("-t", "--targets"): |
| 40 |
# Is the argument a file? |
| 41 |
if os.path.isfile(arg): |
| 42 |
# If the argument is a file, parse it. |
| 43 |
optD['targets'] = list() |
| 44 |
|
| 45 |
try: |
| 46 |
with open(arg, 'r') as fin: |
| 47 |
while True: |
| 48 |
targetsRaw = fin.readlines(8192) |
| 49 |
if not targetsRaw: |
| 50 |
break |
| 51 |
targets = [target.strip(' \t\n,') for target in targetsRaw] |
| 52 |
optD['targets'].extend(targets) |
| 53 |
except IOError as exc: |
| 54 |
type, value, traceback = sys.exc_info() |
| 55 |
print(exc.strerror, file=sys.stderr) |
| 56 |
print('Unable to open ' + "'" + value.filename + "'.", file=sys.stderr) |
| 57 |
istat = bool(1) |
| 58 |
else: |
| 59 |
# Otherwise, parse the argument itself. |
| 60 |
optD['targets'] = arg.split(',') # a list |
| 61 |
elif opt in ("-k", "--keys"): |
| 62 |
# Is the argument a file? |
| 63 |
if os.path.isfile(arg): |
| 64 |
# If the argument is a file, parse it. |
| 65 |
optD['keys'] = list() |
| 66 |
|
| 67 |
try: |
| 68 |
with open(arg, 'r') as fin: |
| 69 |
while True: |
| 70 |
keysRaw = fin.readlines(8192) |
| 71 |
if not keysRaw: |
| 72 |
break |
| 73 |
keys = [key.strip(' \t\n,') for key in keysRaw] |
| 74 |
optD['keys'].extend(keys) |
| 75 |
except IOError as exc: |
| 76 |
type, value, traceback = sys.exc_info() |
| 77 |
print(exc.strerror, file=sys.stderr) |
| 78 |
print('Unable to open ' + "'" + value.filename + "'.", file=sys.stderr) |
| 79 |
istat = bool(1) |
| 80 |
else: |
| 81 |
# Otherwise, parse the argument itself. |
| 82 |
optD['keys'] = arg.split(',') # a comma-separated list |
| 83 |
elif opt in ("-u", "--unique"): |
| 84 |
optD['unique'] = arg.split(',') # a comma-separated list |
| 85 |
elif opt in ("-n", "--dbname"): |
| 86 |
optD['dbname'] = arg |
| 87 |
elif opt in ("-h", "--dbhost"): |
| 88 |
optD['dbhost'] = arg |
| 89 |
elif opt in ("-p", "--dbport"): |
| 90 |
optD['dbport'] = arg |
| 91 |
elif opt == '-d': |
| 92 |
# DoIt! |
| 93 |
optD['doit'] = 1 |
| 94 |
else: |
| 95 |
optD[opt] = arg |
| 96 |
|
| 97 |
if istat or not optD or not 'ns' in optD or not 'targets' in optD or not 'keys' in optD or not 'unique' in optD or not 'dbname' in optD or not 'dbhost' in optD or not 'dbport' in optD: |
| 98 |
print(optD) |
| 99 |
print('Missing required arguments.', file=sys.stderr) |
| 100 |
optD = list() |
| 101 |
return optD |
| 102 |
|
| 103 |
def makeWhere(pkeys, source, target): |
| 104 |
res = '' |
| 105 |
for key in pkeys: |
| 106 |
if len(res) != 0: |
| 107 |
res += ' AND ' |
| 108 |
res += source + '.' + key + ' = ' + target + '.' + key |
| 109 |
return res |
| 110 |
|
| 111 |
rv = RET_SUCCESS |
| 112 |
|
| 113 |
# Parse arguments |
| 114 |
if __name__ == "__main__": |
| 115 |
optD = GetArgs(sys.argv[1:]) |
| 116 |
if not optD: |
| 117 |
rv = RET_INVALIDARG |
| 118 |
else: |
| 119 |
source = optD['ns'] + '.' + optD['table'] |
| 120 |
targets = optD['targets'] |
| 121 |
keys = optD['keys'] |
| 122 |
pkeys = optD['unique'] |
| 123 |
dbuser = pwd.getpwuid(os.getuid())[0] |
| 124 |
dbname = optD['dbname'] |
| 125 |
dbhost = optD['dbhost'] |
| 126 |
dbport = optD['dbport'] |
| 127 |
if 'doit' in optD: |
| 128 |
doit = 1 |
| 129 |
else: |
| 130 |
doit = 0 |
| 131 |
|
| 132 |
if rv == RET_SUCCESS: |
| 133 |
# Connect to the database |
| 134 |
try: |
| 135 |
# The connection is NOT in autocommit mode. If changes need to be saved, then conn.commit() must be called. |
| 136 |
with psycopg2.connect(database=dbname, user=dbuser, host=dbhost, port=dbport) as conn: |
| 137 |
with conn.cursor() as cursor: |
| 138 |
where = makeWhere(pkeys, 'source', 'target') |
| 139 |
for target in targets: |
| 140 |
for key in keys: |
| 141 |
# UPDATE <target> AS target SET NOAA_ARS = source.NOAA_ARS FROM <source> AS source WHERE target.HARPNUM = source.HARPNUM AND target.T_REC_INDEX = source.T_REC_INDEX; |
| 142 |
sql = 'UPDATE ' + target + ' AS target SET ' + key + ' = source.' + key + ' FROM ' + source + ' AS source WHERE ' + where |
| 143 |
|
| 144 |
if doit: |
| 145 |
cursor.execute(sql) |
| 146 |
else: |
| 147 |
print('Executing SQL:\n ==>' + sql) |
| 148 |
|
| 149 |
except psycopg2.Error as exc: |
| 150 |
# Closes the cursor and connection |
| 151 |
print(exc.diag.message_primary, file=sys.stderr) |
| 152 |
# No need to close cursor - leaving the with block does that. |
| 153 |
if not conn: |
| 154 |
rv = RET_DBCONNECT |
| 155 |
else: |
| 156 |
rv = RET_SQL |
| 157 |
|
| 158 |
# There is no need to call conn.commit() since connect() was called from within a with block. If an exception was not raised in the with block, |
| 159 |
# then a conn.commit() was implicitly called. If an exception was raised, then conn.rollback() was implicitly called. |
| 160 |
|
| 161 |
sys.exit(rv) |
| 162 |
|
| 163 |
|
| 164 |
|
| 165 |
|
| 166 |
|
| 167 |
|
