#!/usr/bin/env python 
# ==============================================================================
'''
Usage: gpinitsegment [options]

Options:
  -l, --logdir
  -c, --content
  -h, --host
  -i, --dbid
  -n, --numseg  
  -p, --port  
  -A, --array-name
  -C, --copy 
  -D, --datadir
  -E, --client_encoding
  -G, --gphome
  -K, --checkpoint_segments
  -L, --locale
  -M, --max_connnections
'''

# TODO MIRRORING: NEED TO TAKE in replicationPort

import sys, os, re, optparse
import time, datetime
import signal, subprocess, tempfile

EXECNAME  = os.path.split(__file__)[1]
os_type   = os.uname()[0].lower()
localhost = os.uname()[1]

# ==============================================================================
class SegmentError(Exception): pass


# ==============================================================================
def ParseInput():
    '''
    Parse input
    '''
    try:
        parser = optparse.OptionParser(usage=__doc__, 
                                       add_help_option=False)
        parser.add_option('-?', '--help', action='store_true',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-l', '--logdir', 
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-h', '--host', 
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-A', '--array-name', dest='cluster',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-p', '--port', type='int',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-c', '--content', type='int',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-i', '--dbid', type='int',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-n', '--numseg', type='int',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-D', '--datadir', 
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-C', '--copy', 
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-G', '--gphome', 
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-K', '--checkpoint_segments',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-E', '--client_encoding',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-L', '--locale',
                          help=optparse.SUPPRESS_HELP)
        parser.add_option('-M', '--max_connections',
                          help=optparse.SUPPRESS_HELP)
        (options, args) = parser.parse_args()

        if options.help:
            print __doc__
            sys.exit(0)

        elif len(args) > 0:
            print __doc__
            sys.exit(1)

    except Exception, e:
        print __doc__
        print 'Error parsing input: '
        raise e

    return options

def RunCmd(cmd, env):
    try:
        pipe = subprocess.Popen(cmd, shell=True, env=env,
                                stdout=subprocess.PIPE, 
                                stderr=subprocess.PIPE)
    except OSError, e:
        raise SegmentError(cmd.split(' ')[0] + ' : ' +  str(e))

    try:
        result = pipe.communicate();
    except OSError, e:
        raise SegmentError(cmd + ' : ' +  str(e))

    if (pipe.returncode):
        print result[0].strip()
        print result[1].strip()
        raise SegmentError(cmd + ' : returned ' + str(pipe.returncode))

    return result[0].strip()
    

# ==============================================================================
class GPInitSegment():

    #----------------------------------------------------------------------
    def __init__(self, options=None):
        '''
        Sets the default values for a master segment
        '''

        home = os.environ.get('HOME')
        
        self.logfile  = None
        
        def notNone(a,b):
            if a != None: return a
            return b

        if home == None:
            self.logdir = options.logdir
        else:
            self.logdir = notNone(options.logdir, 
                                  os.path.join(home, 'gpAdminLogs'))        
        if not os.path.exists(self.logdir):
            self.logdir = None

        self.host     = notNone(options.host, localhost)
        self.cluster  = notNone(options.cluster, 'Greenplum')
        self.port     = notNone(options.port, '5432')
        self.content  = notNone(options.content, '-1')
        self.dbid     = notNone(options.dbid, '1')
        self.numseg   = notNone(options.numseg, '-1')
        self.dir      = notNone(options.datadir,  
                                os.environ.get('MASTER_DATA_DIRECTORY'))
        self.gphome   = options.gphome or os.environ.get('GPHOME')
        self.source   = options.copy
        self.encoding = options.client_encoding
        self.checkseg = options.checkpoint_segments
        self.locale   = options.locale
        self.maxconn  = options.max_connections

        self.user     = os.environ.get('USER') or os.environ.get('LOGNAME')

        if not self.gphome:
            print "Could not determine GPHOME"
            sys.exit(0)

        if int(self.numseg) <= 0:
            print "Total number of segments must be > 0"
            sys.exit(1)

        if self.dir[0] != '/':
            self.dir = os.path.join(os.getcwd(), self.dir)

        if self.logdir:
            today = datetime.date.today().strftime('%Y%m%d')
            logname = os.path.join(self.logdir, '%s_%s.log' % (EXECNAME, today))
            self.logfile = open(logname, 'a')            
        

        self.env = {}
        self.env['HOME']    = os.environ.get('HOME')
        self.env['USER']    = self.user
        self.env['LOGNAME'] = self.user
        self.env['GPHOME']  = self.gphome

        # Need to establish an environment to run commands in
        # + Add $GPHOME/bin to the path for this environment
        # + Add /sbin to the path for ifconfig
        path = os.environ.get('PATH') or ''
        path = '%s/bin:%s/ext/python/bin:/sbin:%s' % (self.gphome, self.gphome, path)
        self.env['PATH'] = path

        path = os.environ.get('LD_LIBRARY_PATH') or ''
        path = '%s/lib:%s/ext/python/lib:%s' % (self.gphome, self.gphome, path)
        self.env['LD_LIBRARY_PATH'] = path

        path = os.environ.get('DYLD_LIBRARY_PATH') or ''
        path = '%s/lib:%s/ext/python/lib:%s' % (self.gphome, self.gphome, path)
        self.env['DYLD_LIBRARY_PATH'] = path

    #----------------------------------------------------------------------
    def CreateSegment(self):
        '''
        Creates the new segment via one of three main methods:
           1) Executes gpinitsegment on remote host
           2) Runs initdb in the directory
           3) Copies the data directory from to source location
        '''
        if self.host != localhost:
            self.ExecuteRemote()
        elif self.source:
            self.CopySeg()
        else:
            self.InitSeg()

    #----------------------------------------------------------------------
    def ExecuteRemote(self):
        '''
        Remote execution strategy is simply to use ssh to call 
        gpinitsegment on the remote host. 
        '''

        self.Log('INFO', 'Executing initseg remotely on %s' % self.host)

        ENV      = '%s/bin/lib/gpenv.sh' % self.gphome
        cmd  = 'ssh %s %s %s' % (self.host, ENV, __file__)
        cmd += ' -G %s' % self.gphome
        cmd += ' -A \'\\"%s\\"\'' % self.cluster
        cmd += ' -p %s' % self.port
        cmd += ' -c %s' % self.content
        cmd += ' -i %s' % self.dbid
        cmd += ' -n %s' % self.numseg
        cmd += ' -D %s' % self.dir

        # If using "--copy" then the current hostname will be prefixed 
        # to the copy source if a hostname is not already specified
        if self.source:

            source_re = re.compile('(?:([^:]+)\:(\d*))?(.*)')
            m = source_re.match(self.source)
            if not m:
                raise SegmentError("Invalid source description: '%s'" % self.source)
            s_host = m.group(1)
            s_port = m.group(2)
            s_dir  = m.group(3)

            # make s_dir into a full path
            if s_dir[0] != '/':
                s_dir = os.path.join(os.getcwd(), s_dir)

            if s_host == None:
                s_host = localhost

            cmd += ' -C %s:%s' % (s_host, s_dir)

        result = RunCmd(cmd, self.env)
        if (self.logfile):
            self.logfile.write(result + "\n")
        sys.stdout.write(result + "\n")


    #----------------------------------------------------------------------
    def Log(self, type, msg):
        '''
        Writes a message to stderr and the logfile
        '''
        prefix = ':'.join(
            [time.strftime('%Y%m%d:%H:%M:%S', time.localtime()), EXECNAME, 
             self.host, self.user, '[%s]' % type, '-'])

        for m in msg.splitlines():
            if m.strip() == '': 
                continue

            if (self.logfile):
                self.logfile.write('%s%s\n' % (prefix, m))
                self.logfile.flush()

            sys.stdout.write('%s%s\n' % (prefix, m))

    #----------------------------------------------------------------------
    def MakeDirectory(self):
        '''
        Creates the directory for the segment
        '''
        self.Log('INFO', 'Creating data directory: %s' % self.dir)

        if os.path.exists(self.dir):
            raise SegmentError("'%s' already exists" % self.dir)

        (path, dir) = os.path.split(self.dir)

        if not os.path.exists(path):
            raise SegmentError("'%s' does not exist" % path)

        if not os.access(path, os.W_OK):
            raise SegmentError("'%s' no write permission" % path)

        try:
            os.mkdir(self.dir)
            RunCmd('chmod 700 ' + self.dir, self.env)
        except Exception, e:
            raise SegmentError(str(e))

    #----------------------------------------------------------------------
    def SetPort(self):
        '''
        Modifies postgresql.conf file to have the correct port
        '''

        self.Log('INFO', 'Setting port')
        
        if not os.path.exists(self.dir):
            raise SegmentError("'%s' does not exist" % self.dir)
        pidfile = os.path.join(self.dir, 'postmaster.pid')
        if os.path.exists(pidfile):
            raise SegmentError("Can not adjust port when segment is running")
        
        conf = os.path.join(self.dir, 'postgresql.conf')
        if not os.path.exists(conf):
            raise SegmentError("'%s' does not exist" % conf)

        # Rewrite the postgresql.conf file
        port_re   = re.compile('#?port\s*=\s*\d+.*')
        chkpnt_re = re.compile('#?checkpoint_segments\s*=\s*\d+.*')
        encode_re = re.compile('#?client_encoding\s*=\s*\w+.*')
        old = open(conf, 'r')
        new = open(conf+'.new', 'w')
        for line in old:
            if port_re.match(line):
                new.write('port = %s     # sets the database listener port\n'
                          % self.port)
            elif (self.checkseg and chkpnt_re.match(line)):
                new.write('checkpoint_segments = %s     # in logfile segments, min 1, 16MB each\n'
                          % self.checkseg)
            elif (self.encoding and encode_re.match(line)):
                new.write('client_encoding = %s     # default database encoding\n'
                          % self.encoding)
            else:
                new.write(line)

        # Replace the old file with the new one
        os.rename(conf+'.new', conf)

    def SetHba(self):
        '''
        Modifies pg_hba.conf file to allow access
        '''

        self.Log('INFO', 'Configuring pg_hba.conf')

        # Customize for OS
        if os_type == 'sunos':
            cmd = 'ifconfig -a inet'
        else:
            cmd = 'ifconfig'

        # open pg_hba.conf for writing
        pg_hba = open(os.path.join(self.dir, 'pg_hba.conf'), 'w')

        # Add an entry for local connections by the admin user
        if os_type == 'sunos':
            pg_hba.write('local all all trust\n')
        else:
            pg_hba.write('local all %s ident\n' % self.user)

        # Find all known ip addresses that this host identifies with
        ifconf = RunCmd(cmd, self.env)
        ip_re = re.compile('inet (?:addr:)?(\d+\.\d+\.\d+\.\d+)')
        for line in ifconf.split('\n'):
            m = ip_re.search(line)
            if m:
                ip = m.group(1)
                cidr_suffix = '/128' if ':' in ip else '/32' # MPP-15889
                pg_hba.write('host all all %s%s trust\n' % (ip, cidr_suffix))

        pg_hba.close()


    #----------------------------------------------------------------------
    def SetIdentity(self):
        '''
        Modifies gp_id for the segment in both postgres and template1 dbs
        '''

        ident = "('%s', %s, %s, %s)" % \
            (self.cluster, self.numseg, self.dbid, self.content)
        self.Log('INFO', 'Setting identity %s' % ident)

        options = "PGOPTIONS='-c gp_session_role=utility'"
        psql    = options + ' psql -p %s' % self.port

        def runSQL(sql):
            r = subprocess.call(psql + ' -c "%s" template1' % sql,
                                shell=True,
                                stdout=self.logfile,
                                stderr=self.logfile)
            if r != 0:
                raise SegmentError('Error Executing SQL (return code %s)' % str(r))

        runSQL("delete from gp_id")
        runSQL("insert into gp_id values %s" % ident)
        runSQL("vacuum freeze gp_id")


    #----------------------------------------------------------------------
    def Startup(self):
        '''
        Uses pg_ctl to startup the database:
          + pg_ctl will wait for the startup to complete
          + pg_ctl may gives up after about a minute if the startup fails
        '''

        self.Log('INFO', "Starting Greenplum (port=%s, dir=%s)" 
                 % (self.port, self.dir))

        cmd  = 'pg_ctl start -w -D %s -l %s.log' % (self.dir, self.dir)
        cmd += ' -o "-i -p %s -c gp_role=utility"' % str(self.port)
        self.Log('INFO', cmd)
        retcode = subprocess.call(cmd, shell=True, env=self.env)
        if retcode != 0:
            raise SegmentError('Segment failed to startup, pg_ctl returned %s' 
                               % str(retcode))


    #----------------------------------------------------------------------
    def Shutdown(self):
        '''
        Checks for a postmaster.pid file.
          + If found it calls pg_ctl to stop the database.  
          + pg_ctl will wait until the database is stopped
          + pg_ctl may gives up after about a minute if the shutdown fails
        '''

        if os.path.exists(os.path.join(self.dir, 'postmaster.pid')):
            self.Log('INFO', "Stopping Greenplum (port=%s, dir=%s)" 
                     % (self.port, self.dir))

            cmd  = 'pg_ctl stop -D %s -m fast' % self.dir
            self.Log('INFO', cmd)
            retcode = subprocess.call(cmd, shell=True, env=self.env)
            if retcode != 0:
                raise SegmentError('Segment failed to stop, pg_ctl returned %s' 
                                   % str(retcode))


    #----------------------------------------------------------------------
    def InitSeg(self):
        '''
        
        '''
        self.MakeDirectory()

        self.Log('INFO', 'Running initdb')

        try:
            cmd = 'initdb -D ' + self.dir
            if self.encoding:
                cmd += ' -E ' + self.encoding
            if self.locale:
                cmd += ' --locale=' + self.locale
            if self.maxconn:
                cmd += ' --max_connections=' + self.maxconn

                
            self.Log('INFO', cmd)

            subprocess.call(cmd, shell=True,
                            stdout=self.logfile,
                            stderr=self.logfile)
        except BaseException, e:
            try:     
                os.removedirs(self.dir)
            except:  pass
            raise SegmentError(str(e))

        self.SetPort()
        self.SetHba()
        self.Startup()
        self.SetIdentity()

    #----------------------------------------------------------------------
    def CopySeg(self):
        '''
        '''
        source_re = re.compile('(?:([^:]+)\:(\d*))?(.*)')
        m = source_re.match(self.source)
        if not m:
            raise SegmentError("Invalid source description: '%s'" % self.source)
        s_host = m.group(1)
        s_port = m.group(2)
        s_dir  = m.group(3)

        # make s_dir into a full path
        if s_dir[0] != '/':
            s_dir = os.path.join(os.getcwd(), s_dir)

        # Determine what that host thinks its hostname is
        if not s_host:
            s_host = localhost

        SSH      = 'ssh %s' % s_host
        ENV      = '%s/bin/lib/gpenv.sh' % self.gphome

        HOSTNAME = RunCmd('which hostname', self.env)
        TAR      = RunCmd('which tar', self.env)
        PG_CTL   = RunCmd('which pg_controldata', self.env)
        PG_CTL   = '%s %s' % (ENV, PG_CTL)

        # Validate that it's a real data-directory
        try:
            ctl = RunCmd('%s %s %s' % (SSH, PG_CTL, s_dir), self.env)
        except:
            raise SegmentError("Source Directory '%s' invalid" % s_dir)

        # Everything looks good, get ready to rumble
        self.MakeDirectory()
        
        self.Log('INFO', 'Copying data directory')

        cmd  = "'cd %s && %s -cf - .' | (cd %s && %s -xBf - .)" \
            % (s_dir, TAR, self.dir, TAR)

        RunCmd("%s %s" % (SSH, cmd), self.env)
        self.SetPort()
        self.Startup()
        self.SetIdentity()

        
# ==============================================================================
if __name__ == "__main__":
    try:
        options = ParseInput()

        s = GPInitSegment(options)
        try:
            s.CreateSegment()
        except BaseException, e:
            s.Log('FATAL', str(e))
            raise Exception('')
        finally:
            s.Shutdown()
            
        sys.exit(0)

    except Exception, e:
        print str(e)
        sys.exit(1)


