#!/usr/bin/python
"""
pgtune

Sample usage shown by running with "--help"
"""

import sys
import os
import datetime
import optparse
import csv
import platform

# Windows specific routines
try:
  # ctypes is only available starting in Python 2.5
  from ctypes import *
  # wintypes is only is available on Windows
  from ctypes.wintypes import *
  
  def Win32Memory():
    class memoryInfo(Structure):
        _fields_ = [
            ('dwLength', c_ulong),
            ('dwMemoryLoad', c_ulong),
            ('dwTotalPhys', c_ulong),
            ('dwAvailPhys', c_ulong),
            ('dwTotalPageFile', c_ulong),
            ('dwAvailPageFile', c_ulong),
            ('dwTotalVirtual', c_ulong),
            ('dwAvailVirtual', c_ulong)
        ]
        
    mi = memoryInfo()
    mi.dwLength = sizeof(memoryInfo)
    windll.kernel32.GlobalMemoryStatus(byref(mi))
    return mi.dwTotalPhys

except ImportError:
  # TODO Try and use MFI if we're on Python 2.4 (so no ctypes) but it's available?
  pass

def totalMem():
  try:
    if platform.system()=="Windows":
      totalMem=Win32Memory()
    else:
      # Should work on other, more UNIX-ish platforms
      physPages = os.sysconf("SC_PHYS_PAGES")
      pageSize = os.sysconf("SC_PAGE_SIZE")
      totalMem = physPages * pageSize
    return totalMem
  except:
    return None


class PGConfigLine:
  """
  Stores the value of a single line in the postgresql.conf file, with the 
  following fields:
    lineNumber : integer
    originalLine : string
    commentSection : string
    setsParameter : boolean
  
  If setsParameter is True these will also be set:
    name : string
    readable : string
    raw : string  This is the actual value 
    delimiter (expectations are ' and ")
  """

  def __init__(self,line,num=0):
    self.originalLine=line
    self.lineNumber=num
    self.setsParameter=False

    # Remove comments and edge whitespace
    self.commentSection=""
    commentIndex=line.find('#')
    if commentIndex >= 0:      
      line=line[0:commentIndex]
      self.commentSection=line[commentIndex:]

    line=line.strip()
    if line == "":
      return

    # Split into name,value pair
    equalIndex=line.find('=')
    if equalIndex<0:
      return

    (name,value)=line.split('=')
    name=name.strip()
    value=value.strip()
    self.name=name;
    self.setsParameter=True;

    # Many types of values have ' ' characters around them, strip
    # TODO Set delimiter based on whether there is one here or not
    value=value.rstrip("'")
    value=value.lstrip("'")

    self.readable=value

  def outputFormat(self):
    s=self.originalLine;
    return s

  # Implement a Java-ish interface for this class
  def getName(self):
    return self.name

  def getValue(self):
    return self.readable

  def getLineNumber(self):
    return self.lineNumber

  def isSetting(self):
    return self.setsParameter

  def toString(self):
    s=str(self.lineNumber)+" sets?="+str(self.setsParameter)
    if self.setsParameter:
      s=s+" "+self.getName()+"="+self.getValue()
      # TODO:  Include commentSection, readable,raw, delimiter

    s=s+" originalLine:  "+self.originalLine
    return s


class PGConfigFile:
  """
  Read, write, and manage a postgresql.conf file

  There are two main structures here:

  configFile[]:  Array of PGConfigLine entries for each line in the file
  settingLookup:  Dictionary mapping parameter names to the line that set them
  """

  def __init__(self, filename):
    self.readConfigFile(filename)

  def readConfigFile(self,filename):
    self.filename=filename
    self.settingsLookup={}
    self.configFile=[]

    lineNum=0;
    for line in open(filename):
      line=line.rstrip('\n')
      lineNum=lineNum + 1

      configLine=PGConfigLine(line,lineNum)
      self.configFile.append(configLine)

      if configLine.isSetting():
        self.settingsLookup[configLine.getName()]=configLine

  def updateSetting(self,name,newValue):
    newLineText=str(name)+" = "+str(newValue)+" # pg_generate_conf wizard "+str(datetime.date.today())
    newLine=PGConfigLine(newLineText)

    if self.settingsLookup.has_key(name):
      # Comment out old line
      oldLine=self.settingsLookup[name]
      oldLineNum=oldLine.getLineNumber()
      commentedLineText="# "+oldLine.outputFormat()
      commentedLine=PGConfigLine(commentedLineText,oldLineNum)
      # Subtract one here to adjust for zero offset of array.
      # Any future change that adds lines in-place will need to do something
      # smarter here, because the line numbers won't match the array indexes
      # anymore
      self.configFile[oldLineNum-1]=commentedLine

    self.configFile.append(newLine)
    self.settingsLookup[name]=newLine

  def updateIfLarger(self,name,newValue):
    if self.settingsLookup.has_key(name):
      # TODO This comparison needs all the values converted to numeric form
      # and converted to the same scale before it will work
      if (True):  #newValue > self.settingsLookup[name].getValue():
        self.updateSetting(name,newValue)

  def writeConfigFile(self,fileHandle):
    for l in self.configFile:
      fileHandle.write(l.outputFormat()+"\n")

  def debugPrintInput(self):
    print "Original file:"
    for l in self.configFile:
      print l.toString()

  def debugPrintSettings(self):
    print "Settings listing:"
    for k in self.settingsLookup.keys():
      print k,'=',self.settingsLookup[k].getValue()


class pg_settings:
  """
  Read and index a delimited text dump of a typical pg_settings dump for 
  the appropriate architecture--maximum values are different for some
  settings on 32 and 64 bit platforms.

  The file this needs to operate correctly can be generated with:

  psql postgres -c "COPY (SELECT name,setting,unit,category,short_desc,
  extra_desc,context,vartype,min_val,max_val,enumvals,boot_val 
  FROM pg_settings WHERE NOT source='override') TO '/<path>/pg_settings-<ver>-<bits>'"

  Note that some of these columns (such as boot_val) are only available 
  starting in PostgreSQL 8.4
  """

  def __init__(self):
    self.readConfigFile()

  def readConfigFile(self):
    self.settingsLookup={}
    self.memoryUnits={}

    platformBits=32
    if platform.architecture()[0]=="64bit":  platformBits=64
    # TODO Base this file location on where this script is at
    # TODO Support handling versions other than 8.4
    settingDumpFile="pg_settings-8.4-"+str(platformBits)
    settingColumns=["name","setting","unit","category","short_desc",
      "extra_desc","context","vartype","min_val","max_val","enumvals",
      "boot_val"]
    reader = csv.DictReader(open(settingDumpFile), settingColumns, delimiter="\t")
    for d in reader:
      # Memory units must be specified in some number of kB (never a larger 
      # unit).  Typically they are either "kB" for 1kB or "8kB", unless someone
      # compiled the server with a larger database or xlog block size
      # (BLCKSZ/XLOG_BLCKSZ).  This code has no notion that such a thing is
      # possible though.
      d['memory_unit']=d['unit'].endswith('kB');
      if d['memory_unit']:
        divisor=d['unit'].rstrip('kB')
        if divisor=='':  divisor="1"
        d['memory_divisor']=int(divisor)
      else:
        d['memory_divisor']=None

      self.settingsLookup[d['name']]=d

  def debugPrintSettings(self):
    for key in self.settingsLookup.keys():
      print "key=",key," value=",self.settingsLookup[key]

  def min_val(self,setting):
    return (self.settingsLookup[setting])['min_val']

  def max_val(self,setting):
    return (self.settingsLookup[setting])['max_val']

  def unit(self,setting):
    return (self.settingsLookup[setting])['unit']

  def memory_unit(self,setting):
    return (self.settingsLookup[setting])['memory_unit']

  def vartype(self,setting):
    return (self.settingsLookup[setting])['vartype']

  def show(self,name,value):
    formatted=value
    s=self.settingsLookup[name]
    print >> sys.stderr,"Showing",name,"with value",value

    if s['memory_unit']:
      # Use the same logic as the GUC code that implements "SHOW".  This uses
      # larger units only if there's no loss of resolution in displaying
      # with that value.  Therefore, if using this to output newly assigned
      # values, that value needs to be rounded appropriately if you want
      # it to show up as an even number of MB or GB
      KB_PER_MB=1024
      KB_PER_GB=1024*1024
      if (value % KB_PER_GB == 0):
        value=value/KB_PER_GB
        unit="GB"
      elif (value % KB_PER_MB == 0):
        value=value/KB_PER_MB;
        unit="MB"
      else:
        unit="kB"
      formatted=str(value)+unit

    print >> sys.stderr,"show gives",formatted
    
    return formatted

  def parse_int(self,name,value):
    s=self.settingsLookup[name]
    if self.memory_unit():
      # TODO Finish implementing this so it understands memory units 
      # that have a divisor set
      pass
    return 0

# Beginning of routines for this program

def ReadOptions():
  parser=optparse.OptionParser(
    usage="usage: %prog [options]",
    version="1.0",
    conflict_handler="resolve")
    
  parser.add_option('-i','--input-config',dest="inputConfig",default=None,
    help="Input configuration file")

  parser.add_option('-o','--output-config',dest="outputConfig",default=None, 
    help="Output configuration file, defaults to standard output")
    
  parser.add_option('-M','--memory',dest="totalMemory",default=None, 
    help="Total system memory, will attempt to detect if unspecified")

  parser.add_option('-T','--type',dest="dbType",default="Mixed", 
    help="Database type, defaults to Mixed, valid options are DW, OLTP, Web, Mixed, Desktop")

  parser.add_option('-c','--connections',dest="connections",default=None, 
    help="Maximum number of expected connections, default depends on database type")

  parser.add_option('-D','--debug',action="store_true",dest="debug",
    default="False",help="Enable debugging mode")

  (options,args)=parser.parse_args()
  
  if options.debug==True:
    print "Command line options:  ",options
    print "Command line arguments:  ",args
  
  return (options,args)

def binaryround(value):
  # Keeps the 4 most significant binary bits, truncates the rest so that
  # SHOW will be likely to use a larger divisor
  multiplier=1
  while value>16:
    value=int(value/2)
    multiplier=multiplier * 2
  return multiplier * value

def wizardTune(config,options,settings):
  # We expect the following options are passed into here:
  #
  # dbType:  Defaults to mixed
  # connections:  If missing, will set based on dbType
  # totalMemory:  If missing, will detect

  dbType=options.dbType.lower()

  # Save all settings to be updated as (setting,value) dictionary values
  s={}
  try:
    s['max_connections']={
      'web':200,'oltp':300,'dw':20,'mixed':80,'desktop':5}[dbType]
  except KeyError:
    print "Error:  unexpected setting for dbType"
    os.exit(1)

  # Now that we've screened for that, we know we've got a good dbType and
  # don't have to wrap the rest of these settings in an try block

  # Allow overriding the maximum connections
  if options.connections!=None:
    s['max_connections']=options.connections

  # Estimate memory on this system via parameter or system lookup
  totalMemory=options.totalMemory
  if totalMemory==None:
    totalMemory=totalMem()
  if totalMemory==None:
    print "Error:  total memory not specified and unable to detect"
    os.exit(1)

  kb=1024;
  mb=1024*kb
  gb=1024*mb

  # Memory allocation
  # Extract some values just to make the code below more compact
  # The base unit for memory types is the kB, so scale system memory to that
  mem=totalMemory / kb
  con=s['max_connections']

  if totalMemory>=(256*mb):
    s['shared_buffers']={
      'web':mem/4, 'oltp':mem/4,'dw':mem/4,
      'mixed':mem/4, 'desktop':mem/16}[dbType]

    s['effective_cache_size']={
      'web':mem*3/4, 'oltp':mem*3/4,'dw':mem*3/4,
      'mixed':mem*3/4,'desktop':mem/4}[dbType]

    s['work_mem']={
      'web':mem/con, 'oltp':mem/con,'dw':mem/con/2,
      'mixed':mem/con/2,'desktop':mem/con/6}[dbType]

    s['maintenance_work_mem']={
      'web':mem/16, 'oltp':mem/16,'dw':mem/8,
      'mixed':mem/16,'desktop':mem/16}[dbType]
    # Cap maintenence RAM at 1GB on servers with lots of memory
    # (Remember that the setting is in terms of kB here)
    if s['maintenance_work_mem']>(1*mb):
      s['maintenance_work_mem']=1*mb;

  else:
    # TODO HINT about this tool not being optimal for low memory systems
    pass

  # Checkpoint parameters
  s['checkpoint_segments']={
    'web':8, 'oltp':16, 'dw':64,
    'mixed':16, 'desktop':3}[dbType]

  s['checkpoint_completion_target']={
    'web':0.7, 'oltp':0.9, 'dw':0.9,
    'mixed':0.9, 'desktop':0.5}[dbType]

  s['wal_buffers']=512 * s['checkpoint_segments']

  # Paritioning and statistics
  s['constraint_exclusion']={
    'web':'off', 'oltp':'off', 'dw':'on', 
    'mixed':'on', 'desktop':'off'}[dbType]

  s['default_statistics_target']={
    'web':10, 'oltp':10, 'dw':100, 
    'mixed':50, 'desktop':10}[dbType]

  for k in s.keys():
    value=s[k]
    if settings.memory_unit(k):
      value=binaryround(s[k])
    # TODO Pass a parameter to updateSetting that has it throw a HINT if you're reducing a value
    config.updateSetting(k,settings.show(k,value))

if __name__=='__main__':
  (options,args)=ReadOptions() 
  
  configFile=options.inputConfig
  if configFile==None:
    print >> sys.stderr,"Can't do anything without an input config file; try --help"
    sys.exit(1)
    # TODO Show usage here
    
  config=PGConfigFile(configFile)

  if (options.debug==True):  
    config.debugPrintInput()
    print
    config.debugPrintSettings()

  settings=pg_settings()

  wizardTune(config,options,settings)
  
  outputFileName=options.outputConfig
  if outputFileName==None:  
    outputFile=sys.stdout
  else:
    outputFile=open(outputFileName,'w')

  config.writeConfigFile(outputFile)
