tools/sbin/gpfixuserlimts

#!/usr/bin/env 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.
''' 
USAGE:  gpfixuserlimts -f|--file FILE -u|--user USER
        where FILE is the file to be modified
        where USER is the user for whom limits are being checked and set
'''

import os, sys, re, shutil

try:
    from optparse import Option, OptionParser 
    from gppylib.gpparseopts import OptParser, OptChecker
except ImportError, e:    
    sys.exit('Cannot import modules.  Please check that you have sourced greenplum_path.sh.  Detail: ' + str(e))

options = None
tmpfile = None
bakfile = None
   
# True is an error
def parseargs():

    global options
    global tmpfile
    global bakfile

    parser = OptParser(option_class=OptChecker)
    parser.remove_option('-h')
    parser.add_option('-h', '-?', '--help', action='store_true')
    parser.add_option('-f', '--file', type='string')
    parser.add_option('-u', '--user', type='string')
    (options, args) = parser.parse_args()

    global gphome
    gphome = os.environ.get('GPHOME')
    if not gphome:
        sys.stderr.write("GPHOME not set")
        return True

    if options.help:
        sys.stderr.write(__doc__)
        sys.exit(0)

    if not options.file:
        sys.stderr.write("--file must be specified")
        return True

    tmpfile = "%s.tmp" % options.file
    bakfile = "%s.bak" % options.file

    if not os.path.exists(options.file):
        sys.stderr.write("file %s does not exist" % options.file)
        return True

    if not options.user:
        sys.stderr.write("--user must be specified")
        return True

    if options.user == "root":
        sys.stderr.write("root is not a valid user to be specified to this script")
        return True

    return False

if (parseargs()):
    sys.exit(1)

class Item:
    def __init__(self, itemid, field1, field2, minval):
        self.itemid = itemid
        self.field1 = field1
        self.field2 = field2
        self.minval = int(minval)
        self.satisfied = False
    
    def debugStr(self):
        return "%s %s %s satisfied=%s" % (self.field1, self.field2, self.minval, self.satisfied)

items = {}
items[1] = Item(1, 'soft', 'nofile', 65536)
items[2] = Item(2, 'hard', 'nofile', 65536)
items[3] = Item(3, 'soft', 'nproc', 131072)
items[4] = Item(4, 'hard', 'nproc', 131072)

linenumbersToRemove = set()

########################################################
## GATHER INFO
########################################################
try:
    fd = open(options.file, "r")
    linenumber = 0
    for line in fd:
        linenumber = linenumber + 1
        fields = line.strip().split()
        if len(fields) != 4:
            continue

        wildcard = False
        user = False
        index = None

        if fields[0] == "*":
            wildcard = True

        if fields[0] == options.user:
            user = True

        if not wildcard and not user:
            continue

        for k in items.keys():
            if fields[1] != items[k].field1:
                continue
            if fields[2] != items[k].field2:
                continue
            index = k
            break

        if not index:
            continue

        # we have a match
        try:
            fields[3] = int(fields[3])
        except:
            fields[3] = 0

        if wildcard: 

            if  items[index].minval <= fields[3]:
                # WILDCARD RULE IS GOOD
                items[index].satisfied = True

        else: # matches our user

            if  items[index].minval <= fields[3]:
                # USER RULE IS GOOD
                items[index].satisfied = True

            else:
                # USER RULE THAT IS NOT GOOD
                # REMOVE THIS LINE FROM FILE
                linenumbersToRemove.add(linenumber)
    fd.close()
except Exception,e:
    sys.stderr.write("Error parsing file %s\n" % e.__str__())
    sys.exit(1)

########################################################
## WRITE NEW FILE
########################################################
try:
    output = open(tmpfile, "w")
    input = open(options.file, "r")
    linenumber = 0
    for line in input:
        linenumber = linenumber + 1

        if linenumber in linenumbersToRemove:
            output.write("#%s" % line)
        else:
            output.write("%s" % line)

    # new lines
    for k in items.keys():
        if items[k].satisfied:
            continue
        output.write("%s %s %s %s\n" % (options.user, items[k].field1, items[k].field2, items[k].minval))

    output.close()
    input.close()
except Exception,e:
    sys.stderr.write("Error writing file %s\n" % e.__str__())
    sys.exit(1)
 
########################################################
## SWAP  FILES
########################################################

try:
    shutil.copyfile(options.file, bakfile)
    shutil.move(tmpfile, options.file)
except Exception,e:
    sys.stderr.write("Error swapping files %s\n" % e.__str__())
    sys.exit(1)

sys.exit(0)