WIP-pyshapelib-bramz/test/postgissupport.py

# Copyright (C) 2003 by Intevation GmbH
# Authors:
# Bernhard Herzog <[email protected]>
#
# This program is free software under the GPL (>=v2)
# Read the file COPYING coming with the software for details.

"""Support module for tests that use a live PostGIS database"""

__version__ = "$Revision$"
# $Source$
# $Id$

import sys
import os
import time
import popen2
import shutil
import traceback

import support

try:
    import psycopg
except ImportError:
    psycopg = None

#
#       Helper code
#

def run_config_script(cmdline):
    """Run command cmdline and return its stdout or none in case of errors"""
    pipe = os.popen(cmdline)
    result = pipe.read()
    if pipe.close() is not None:
        raise RuntimeError('Command %r failed' % cmdline)
    return result

def run_command(command, outfilename = None):
    """Run command as a subprocess and send its stdout and stderr to outfile

    The subprocess is run synchroneously so the function returns once
    the subprocess has termninated. If the process' exit code is not
    zero raise a RuntimeError.

    If outfilename is None stdout and stderr are still captured but they
    are ignored and not written to any file.
    """
    proc = popen2.Popen4(command)
    proc.tochild.close()
    output = proc.fromchild.read()
    status = proc.wait()
    if outfilename is not None:
        outfile = open(outfilename, "w")
        outfile.write(output)
        outfile.close()
    if not os.WIFEXITED(status) or os.WEXITSTATUS(status) != 0:
        if outfilename:
            message = "see %s" % outfilename
        else:
            message = output
        raise RuntimeError("command %r exited with code %d.\n%s"
                           % (command, status, message))


def run_boolean_command(command):
    """
    Run command as a subprocess silently and return whether it ran successfully

    Silently means that all output is captured and ignored. The exit
    status is true if the command ran successfull, i.e. it terminated by
    exiting and returned as zero exit code and false other wise
    """
    try:
        run_command(command, None)
        return 1
    except RuntimeError:
        pass
    return 0


#
#       PostgreSQL and database
#

class PostgreSQLServer:

    """A PostgreSQL server

    Instances of this class represent a PostgreSQL server with postgis
    extensions run explicitly for the test cases. Such a server has its
    own database directory and its own directory for the unix sockets so
    that it doesn't interfere with any other PostgreSQL server already
    running on the system.
    """

    def __init__(self, dbdir, port, postgis_sql, socket_dir):
        """Initialize the PostgreSQLServer object

        Parameters:

          dbdir -- The directory for the databases
          port -- The port to use
          postgis_sql -- The name of the file with the SQL statements to
                         initialize a database for postgis.
          socket_dir -- The directory for the socket files.

        When connecting to the database server use the port and host
        instance variables.
        """
        self.dbdir = dbdir
        self.port = port
        self.postgis_sql = postgis_sql
        self.socket_dir = socket_dir

        # For the client side the socket directory can be used as the
        # host if the name starts with a slash.
        self.host = os.path.abspath(socket_dir)

        # name and password for the admin and an unprivileged user
        self.admin_name = "postgres"
        self.admin_password = "postgres"
        self.user_name = "observer"
        self.user_password = "telescope"

        # Map db names to db objects
        self.known_dbs = {}

    def createdb(self):
        """Create the database in dbdir and start the server.

        First check whether the dbdir already exists and if necessary
        stop an already running postmaster and remove the dbdir
        directory completely. Then create a new database cluster in the
        dbdir and start a postmaster.
        """
        if os.path.isdir(self.dbdir):
            if self.is_running():
                self.shutdown()
            shutil.rmtree(self.dbdir)
        os.mkdir(self.dbdir)

        run_command(["initdb", "-D", self.dbdir, "-U", self.admin_name],
                    os.path.join(self.dbdir, "initdb.log"))

        extra_opts = "-p %d" % self.port
        if self.socket_dir is not None:
            extra_opts += " -k %s" % self.socket_dir
        run_command(["pg_ctl", "-D", self.dbdir,
                     "-l", os.path.join(self.dbdir, "logfile"),
                     "-o", extra_opts, "start"],
                    os.path.join(self.dbdir, "pg_ctl-start.log"))
        # the -w option of pg_ctl doesn't work properly when the port is
        # not the default port, so we have to implement waiting for the
        # server ourselves
        self.wait_for_postmaster()

        self.alter_user(self.admin_name, self.admin_password)
        self.create_user(self.user_name, self.user_password)

    def wait_for_postmaster(self):
        """Return when the database server is running

        Internal method to wait until the postmaster process has been
        started and is ready for client connections.
        """
        max_count = 60
        count = 0
        while count < max_count:
            try:
                run_command(["psql", "-l", "-p", str(self.port),
                             "-h", self.host, "-U", self.admin_name],
                            os.path.join(self.dbdir, "psql-%d.log" % count))
            except RuntimeError:
                pass
            except:
                traceback.print_exc()
            else:
                break
            time.sleep(0.5)
            count += 1
        else:
            raise RuntimeError("postmaster didn't start")

    def is_running(self):
        """Return true a postmaster process is running on self.dbdir

        This method runs pg_ctl status on the dbdir so even if the
        object has just been created it is possible that this method
        returns true if there's still a postmaster process running for
        self.dbdir.
        """
        return run_boolean_command(["pg_ctl", "-D", self.dbdir, "status"])

    def shutdown(self):
        """Stop the postmaster running for self.dbdir"""
        run_command(["pg_ctl", "-m", "fast", "-D", self.dbdir, "stop"],
                    os.path.join(self.dbdir, "pg_ctl-stop.log"))

    def new_postgis_db(self, dbname, tables = None):
        """Create and return a new PostGISDatabase object using self as server
        """
        db = PostGISDatabase(self, self.postgis_sql, dbname, tables = tables)
        db.initdb()
        self.known_dbs[dbname] = db
        return db

    def get_static_data_db(self, dbname, tables):
        """Return a PostGISDatabase for a database with the given static data

        If no databasse of the name dbname exists, create a new one via
        new_postgis_db and upload the data.

        If a database of the name dbname already exists and uses the
        indicated data, return that. If the already existing db uses
        different data raise a value error.

        The tables argument should be a sequence of table specifications
        where each specifications is a (tablename, shapefilename) pair.
        """
        db = self.known_dbs.get(dbname)
        if db is not None:
            if db.has_data(tables):
                return db
            raise ValueError("PostGISDatabase named %r doesn't have tables %r"
                             % (dbname, tables))
        return self.new_postgis_db(dbname, tables)

    def get_default_static_data_db(self):
        dbname = "PostGISStaticTests"
        tables = [
            # Direct copies of the shapefiles. The shapeids are exactly
            # the same.
            ("landmarks", os.path.join("..", "Data", "iceland",
                                       "cultural_landmark-point.shp"),
             [("gid_offset", 1000)]),
            ("political", os.path.join("..", "Data", "iceland",
                                             "political.shp")),
            ("roads", os.path.join("..", "Data", "iceland",
                                         "roads-line.shp")),

            # The polygon data as a MULTIPOLYGON geometry type
            ("political_multi", os.path.join("..", "Data", "iceland",
                                             "political.shp"),
             [("force_wkt_type", "MULTIPOLYGON")]),
            ]
        return self.get_static_data_db(dbname, tables)

    def connection_params(self, user):
        """Return the connection parameters for the given user

        The return value is a dictionary suitable as keyword argument
        list to PostGISConnection. The user parameter may be either
        'admin' to connect as admin or 'user' to connect as an
        unprivileged user.
        """
        return {"host": self.host, "port": self.port,
                "user": getattr(self, user + "_name"),
                "password": getattr(self, user + "_password")}

    def connection_string(self, user):
        """Return (part of) the connection string to pass to psycopg.connect

        The string contains host, port, user and password. The user
        parameter must be either 'admin' or 'user', as for
        connection_params.
        """
        params = []
        for key, value in self.connection_params(user).items():
            # FIXME: this doesn't do quiting correctly but that
            # shouldn't be much of a problem (people shouldn't be using
            # single quotes in filenames anyway :) )
            params.append("%s='%s'" % (key, value))
        return " ".join(params)

    def execute_sql(self, dbname, user, sql):
        """Execute the sql statament

        The user parameter us used as in connection_params. The dbname
        parameter must be the name of a database in the cluster.
        """
        conn = psycopg.connect("dbname=%s " % dbname
                               + self.connection_string(user))
        cursor = conn.cursor()
        cursor.execute(sql)
        conn.commit()
        conn.close()

    def require_authentication(self, required):
        """Switch authentication requirements on or off

        When started for the first time no passwords are required. Some
        tests want to explicitly test whether Thuban's password
        infrastructure works and switch password authentication on
        explicitly. When switching it on, there should be a
        corresponding call to switch it off again in the test case'
        tearDown method or in a finally: block.
        """
        if required:
            contents = "local all password\n"
        else:
            contents = "local all trust\n"
        f = open(os.path.join(self.dbdir, "pg_hba.conf"), "w")
        f.write(contents)
        f.close()
        run_command(["pg_ctl", "-D", self.dbdir, "reload"],
                    os.path.join(self.dbdir, "pg_ctl-reload.log"))


    def create_user(self, username, password):
        """Create user username with password in the database"""
        self.execute_sql("template1", "admin",
                         "CREATE USER %s PASSWORD '%s';" % (username,password))

    def alter_user(self, username, password):
        """Change the user username's password in the database"""
        self.execute_sql("template1", "admin",
                         "ALTER USER %s PASSWORD '%s';" % (username,password))


class PostGISDatabase:

    """A PostGIS database in a PostgreSQLServer"""

    def __init__(self, server, postgis_sql, dbname, tables = None):
        """Initialize the PostGISDatabase

        Parameters:

            server -- The PostgreSQLServer instance containing the
                database

            postgis_sql -- Filename of the postgis.sql file with the
                postgis initialization code

            dbname -- The name of the database

            tables -- Optional description of tables to create in the
                new database. If given it should be a list of
                (tablename, shapefilename) pairs meaning that a table
                tablename will be created with the contents of the given
                shapefile or (tablename, shapefilename, extraargs)
                triples. The extraargs should be a list of key, value
                pairs to use as keyword arguments to upload_shapefile.
        """
        self.server = server
        self.postgis_sql = postgis_sql
        self.dbname = dbname
        self.tables = tables

    def initdb(self):
        """Remove the old db directory and create and initialize a new database
        """
        run_command(["createdb", "-p", str(self.server.port),
                     "-h", self.server.host, "-U", self.server.admin_name,
                     self.dbname],
                    os.path.join(self.server.dbdir, "createdb.log"))
        run_command(["createlang", "-p", str(self.server.port),
                     "-h", self.server.host,  "-U", self.server.admin_name,
                     "plpgsql", self.dbname],
                    os.path.join(self.server.dbdir, "createlang.log"))
        # for some reason psql doesn't exit with an error code if the
        # file given as -f doesn't exist, so we check manually by trying
        # to open it before we run psql
        f = open(self.postgis_sql)
        f.close()
        del f
        run_command(["psql", "-f", self.postgis_sql, "-d", self.dbname,
                     "-p", str(self.server.port), "-h", self.server.host,
                     "-U", self.server.admin_name],
                     os.path.join(self.server.dbdir, "psql.log"))

        self.server.execute_sql(self.dbname, "admin",
                                "GRANT SELECT ON geometry_columns TO PUBLIC;")

        if self.tables is not None:
            def unpack(item):
                extra = {"force_wkt_type": None, "gid_offset": 0}
                if len(info) == 2:
                    tablename, shapefile = info
                else:
                    tablename, shapefile, kw = info
                    for key, val in kw:
                        extra[key] = val
                return tablename, shapefile, extra

            for info in self.tables:
                tablename, shapefile, kw = unpack(info)
                upload_shapefile(shapefile, self, tablename, **kw)

    def has_data(self, tables):
        return self.tables == tables


def find_postgis_sql():
    """Return the name of the postgis_sql file

    A postgis installation usually has the postgis_sql file in
    PostgreSQL's datadir (i.e. the directory where PostgreSQL keeps
    static files, not the directory containing the databases).
    Unfortunately there's no way to determine the name of this directory
    with pg_config so we assume here that it's
    $bindir/../share/postgresql/.
    """
    bindir = run_config_script("pg_config --bindir").strip()
    return os.path.join(bindir, "..", "share", "postgresql",
                        "contrib", "postgis.sql")

_postgres_server = None
def get_test_server():
    """Return the test database server object.

    If it doesn't exist yet, create it first.

    The server will use the directory postgis under the temp dir (as
    defined by support.create_temp_dir()) for the database cluster.
    Sockets will be created in tempdir.
    """
    global _postgres_server
    if _postgres_server is None:
        tempdir = support.create_temp_dir()
        dbdir = os.path.join(tempdir, "postgis")
        socket_dir = tempdir

        _postgres_server = PostgreSQLServer(dbdir, 6543, find_postgis_sql(),
                                            socket_dir = socket_dir)
        _postgres_server.createdb()

    return _postgres_server

def shutdown_test_server():
    """Shutdown the test server if it is running"""
    global _postgres_server
    if _postgres_server is not None:
        _postgres_server.shutdown()
        _postgres_server = None


def reason_for_not_running_tests():
    """
    Determine whether postgis tests can be run and return a reason they can't

    There's no fool-proof way to reliably determine this short of
    actually running the tests but we try the following here:

     - test whether pg_ctl --help can be run successfully
     - test whether the postgis_sql can be opened
       The name of the postgis_sql file is determined by find_postgis_sql()
     - psycopg can be imported successfully.
    """
    try:
        run_command(["pg_ctl", "--help"], None)
    except RuntimeError:
        return "Can't run PostGIS tests because pg_ctl fails"

    try:
        postgis_sql = find_postgis_sql()
    except:
        return "Can't run PostGIS tests because postgis.sql can't be found"

    try:
        f = open(postgis_sql)
        f.close()
    except:
        return "Can't run PostGIS tests because postgis.sql can't be opened"

    # The test for psycopg was already done when this module was
    # imported so we only have to check whether it was successful
    if psycopg is None:
        return "Can't run PostGIS tests because psycopg can't be imported"

    return ""


_cannot_run_postgis_tests = None
def skip_if_no_postgis():
    global _cannot_run_postgis_tests
    if _cannot_run_postgis_tests is None:
        _cannot_run_postgis_tests = reason_for_not_running_tests()
    if _cannot_run_postgis_tests:
        raise support.SkipTest(_cannot_run_postgis_tests)

def coords_to_point(coords):
    """Return string with a WKT representation of the point in coords"""
    x, y = coords[0]
    return "POINT(%r %r)" % (x, y)

def coords_to_polygon(coords):
    """Return string with a WKT representation of the polygon in coords"""
    poly = []
    for ring in coords:
        poly.append(", ".join(["%r %r" % p for p in ring]))
    return "POLYGON((%s))" % "), (".join(poly)

def coords_to_multilinestring(coords):
    """Return string with a WKT representation of the arc in coords"""
    poly = []
    for ring in coords:
        poly.append(", ".join(["%r %r" % p for p in ring]))
    return "MULTILINESTRING((%s))" % "), (".join(poly)

def coords_to_multipolygon(coords):
    """Return string with a WKT representation of the polygon in coords"""
    poly = []
    for ring in coords:
        poly.append(", ".join(["%r %r" % p for p in ring]))
    return "MULTIPOLYGON(((%s)))" % ")), ((".join(poly)

wkt_converter = {
    "POINT": coords_to_point,
    "MULTILINESTRING": coords_to_multilinestring,
    "POLYGON": coords_to_polygon,
    "MULTIPOLYGON": coords_to_multipolygon,
    }

def upload_shapefile(filename, db, tablename, force_wkt_type = None,
                     gid_offset = 0):
    """Upload a shapefile into a new database table

    Parameters:

    filename -- The name of the shapefile

    db -- The PostGISDatabase instance representing the database

    tablename -- The name of the table to create and into which the data
                is to be inserted

    force_wkt_type -- If given the real WKT geometry type to use instead
                of the default that would be chosen based on the type of
                the shapefile

    gid_offset -- A number to add to the shapeid to get the value for
                the gid column (default 0)
    """
    import dbflib, shapelib

    # We build this map here because we need shapelib which can only be
    # imported after support.initthuban has been called which we can't
    # easily do in this module because it's imported by support.
    shp_to_wkt = {
        shapelib.SHPT_POINT: "POINT",
        shapelib.SHPT_ARC: "MULTILINESTRING",
        shapelib.SHPT_POLYGON: "POLYGON",
        }

    server = db.server
    dbname = db.dbname
    conn = psycopg.connect("dbname=%s " % dbname
                           + db.server.connection_string("admin"))
    cursor = conn.cursor()

    shp = shapelib.ShapeFile(filename)
    dbf = dbflib.DBFFile(filename)
    typemap = {dbflib.FTString: "VARCHAR",
               dbflib.FTInteger: "INTEGER",
               dbflib.FTDouble: "DOUBLE PRECISION"}

    insert_formats = ["%(gid)s"]
    fields = ["gid INT"]
    for i in range(dbf.field_count()):
        ftype, name, width, prec = dbf.field_info(i)
        fields.append("%s %s" % (name, typemap[ftype]))
        insert_formats.append("%%(%s)s" % name)
    stmt = "CREATE TABLE %s (\n    %s\n);" % (tablename,
                                              ",\n    ".join(fields))
    cursor.execute(stmt)
    #print stmt

    numshapes, shapetype, mins, maxs = shp.info()
    wkttype =  shp_to_wkt[shapetype]
    if force_wkt_type:
        wkttype = force_wkt_type
    convert = wkt_converter[wkttype]

    cursor.execute("select AddGeometryColumn('%(dbname)s',"
                   "'%(tablename)s', 'the_geom', '-1', '%(wkttype)s', 2);"
                   % locals())

    insert_formats.append("GeometryFromText(%(the_geom)s, -1)")

    insert = ("INSERT INTO %s VALUES (%s)"
              % (tablename, ", ".join(insert_formats)))

    for i in range(numshapes):
        data = dbf.read_record(i)
        data["tablename"] = tablename
        data["gid"] = i + gid_offset
        data["the_geom"] = convert(shp.read_object(i).vertices())
        #print insert % data
        cursor.execute(insert, data)

    cursor.execute("GRANT SELECT ON %s TO PUBLIC;" % tablename)

    conn.commit()
1	# Copyright (C) 2003 by Intevation GmbH
2	# Authors:
3	# Bernhard Herzog <[email protected]>
4	#
5	# This program is free software under the GPL (>=v2)
6	# Read the file COPYING coming with the software for details.
7
8	"""Support module for tests that use a live PostGIS database"""
9
10	__version__ = "$Revision$"
11	# $Source$
12	# $Id$
13
14	import sys
15	import os
16	import time
17	import popen2
18	import shutil
19	import traceback
20
21	import support
22
23	try:
24	import psycopg
25	except ImportError:
26	psycopg = None
27
28	#
29	# Helper code
30	#
31
32	def run_config_script(cmdline):
33	"""Run command cmdline and return its stdout or none in case of errors"""
34	pipe = os.popen(cmdline)
35	result = pipe.read()
36	if pipe.close() is not None:
37	raise RuntimeError('Command %r failed' % cmdline)
38	return result
39
40	def run_command(command, outfilename = None):
41	"""Run command as a subprocess and send its stdout and stderr to outfile
42
43	The subprocess is run synchroneously so the function returns once
44	the subprocess has termninated. If the process' exit code is not
45	zero raise a RuntimeError.
46
47	If outfilename is None stdout and stderr are still captured but they
48	are ignored and not written to any file.
49	"""
50	proc = popen2.Popen4(command)
51	proc.tochild.close()
52	output = proc.fromchild.read()
53	status = proc.wait()
54	if outfilename is not None:
55	outfile = open(outfilename, "w")
56	outfile.write(output)
57	outfile.close()
58	if not os.WIFEXITED(status) or os.WEXITSTATUS(status) != 0:
59	if outfilename:
60	message = "see %s" % outfilename
61	else:
62	message = output
63	raise RuntimeError("command %r exited with code %d.\n%s"
64	% (command, status, message))
65
66
67	def run_boolean_command(command):
68	"""
69	Run command as a subprocess silently and return whether it ran successfully
70
71	Silently means that all output is captured and ignored. The exit
72	status is true if the command ran successfull, i.e. it terminated by
73	exiting and returned as zero exit code and false other wise
74	"""
75	try:
76	run_command(command, None)
77	return 1
78	except RuntimeError:
79	pass
80	return 0
81
82
83	#
84	# PostgreSQL and database
85	#
86
87	class PostgreSQLServer:
88
89	"""A PostgreSQL server
90
91	Instances of this class represent a PostgreSQL server with postgis
92	extensions run explicitly for the test cases. Such a server has its
93	own database directory and its own directory for the unix sockets so
94	that it doesn't interfere with any other PostgreSQL server already
95	running on the system.
96	"""
97
98	def __init__(self, dbdir, port, postgis_sql, socket_dir):
99	"""Initialize the PostgreSQLServer object
100
101	Parameters:
102
103	dbdir -- The directory for the databases
104	port -- The port to use
105	postgis_sql -- The name of the file with the SQL statements to
106	initialize a database for postgis.
107	socket_dir -- The directory for the socket files.
108
109	When connecting to the database server use the port and host
110	instance variables.
111	"""
112	self.dbdir = dbdir
113	self.port = port
114	self.postgis_sql = postgis_sql
115	self.socket_dir = socket_dir
116
117	# For the client side the socket directory can be used as the
118	# host if the name starts with a slash.
119	self.host = os.path.abspath(socket_dir)
120
121	# name and password for the admin and an unprivileged user
122	self.admin_name = "postgres"
123	self.admin_password = "postgres"
124	self.user_name = "observer"
125	self.user_password = "telescope"
126
127	# Map db names to db objects
128	self.known_dbs = {}
129
130	def createdb(self):
131	"""Create the database in dbdir and start the server.
132
133	First check whether the dbdir already exists and if necessary
134	stop an already running postmaster and remove the dbdir
135	directory completely. Then create a new database cluster in the
136	dbdir and start a postmaster.
137	"""
138	if os.path.isdir(self.dbdir):
139	if self.is_running():
140	self.shutdown()
141	shutil.rmtree(self.dbdir)
142	os.mkdir(self.dbdir)
143
144	run_command(["initdb", "-D", self.dbdir, "-U", self.admin_name],
145	os.path.join(self.dbdir, "initdb.log"))
146
147	extra_opts = "-p %d" % self.port
148	if self.socket_dir is not None:
149	extra_opts += " -k %s" % self.socket_dir
150	run_command(["pg_ctl", "-D", self.dbdir,
151	"-l", os.path.join(self.dbdir, "logfile"),
152	"-o", extra_opts, "start"],
153	os.path.join(self.dbdir, "pg_ctl-start.log"))
154	# the -w option of pg_ctl doesn't work properly when the port is
155	# not the default port, so we have to implement waiting for the
156	# server ourselves
157	self.wait_for_postmaster()
158
159	self.alter_user(self.admin_name, self.admin_password)
160	self.create_user(self.user_name, self.user_password)
161
162	def wait_for_postmaster(self):
163	"""Return when the database server is running
164
165	Internal method to wait until the postmaster process has been
166	started and is ready for client connections.
167	"""
168	max_count = 60
169	count = 0
170	while count < max_count:
171	try:
172	run_command(["psql", "-l", "-p", str(self.port),
173	"-h", self.host, "-U", self.admin_name],
174	os.path.join(self.dbdir, "psql-%d.log" % count))
175	except RuntimeError:
176	pass
177	except:
178	traceback.print_exc()
179	else:
180	break
181	time.sleep(0.5)
182	count += 1
183	else:
184	raise RuntimeError("postmaster didn't start")
185
186	def is_running(self):
187	"""Return true a postmaster process is running on self.dbdir
188
189	This method runs pg_ctl status on the dbdir so even if the
190	object has just been created it is possible that this method
191	returns true if there's still a postmaster process running for
192	self.dbdir.
193	"""
194	return run_boolean_command(["pg_ctl", "-D", self.dbdir, "status"])
195
196	def shutdown(self):
197	"""Stop the postmaster running for self.dbdir"""
198	run_command(["pg_ctl", "-m", "fast", "-D", self.dbdir, "stop"],
199	os.path.join(self.dbdir, "pg_ctl-stop.log"))
200
201	def new_postgis_db(self, dbname, tables = None):
202	"""Create and return a new PostGISDatabase object using self as server
203	"""
204	db = PostGISDatabase(self, self.postgis_sql, dbname, tables = tables)
205	db.initdb()
206	self.known_dbs[dbname] = db
207	return db
208
209	def get_static_data_db(self, dbname, tables):
210	"""Return a PostGISDatabase for a database with the given static data
211
212	If no databasse of the name dbname exists, create a new one via
213	new_postgis_db and upload the data.
214
215	If a database of the name dbname already exists and uses the
216	indicated data, return that. If the already existing db uses
217	different data raise a value error.
218
219	The tables argument should be a sequence of table specifications
220	where each specifications is a (tablename, shapefilename) pair.
221	"""
222	db = self.known_dbs.get(dbname)
223	if db is not None:
224	if db.has_data(tables):
225	return db
226	raise ValueError("PostGISDatabase named %r doesn't have tables %r"
227	% (dbname, tables))
228	return self.new_postgis_db(dbname, tables)
229
230	def get_default_static_data_db(self):
231	dbname = "PostGISStaticTests"
232	tables = [
233	# Direct copies of the shapefiles. The shapeids are exactly
234	# the same.
235	("landmarks", os.path.join("..", "Data", "iceland",
236	"cultural_landmark-point.shp"),
237	[("gid_offset", 1000)]),
238	("political", os.path.join("..", "Data", "iceland",
239	"political.shp")),
240	("roads", os.path.join("..", "Data", "iceland",
241	"roads-line.shp")),
242
243	# The polygon data as a MULTIPOLYGON geometry type
244	("political_multi", os.path.join("..", "Data", "iceland",
245	"political.shp"),
246	[("force_wkt_type", "MULTIPOLYGON")]),
247	]
248	return self.get_static_data_db(dbname, tables)
249
250	def connection_params(self, user):
251	"""Return the connection parameters for the given user
252
253	The return value is a dictionary suitable as keyword argument
254	list to PostGISConnection. The user parameter may be either
255	'admin' to connect as admin or 'user' to connect as an
256	unprivileged user.
257	"""
258	return {"host": self.host, "port": self.port,
259	"user": getattr(self, user + "_name"),
260	"password": getattr(self, user + "_password")}
261
262	def connection_string(self, user):
263	"""Return (part of) the connection string to pass to psycopg.connect
264
265	The string contains host, port, user and password. The user
266	parameter must be either 'admin' or 'user', as for
267	connection_params.
268	"""
269	params = []
270	for key, value in self.connection_params(user).items():
271	# FIXME: this doesn't do quiting correctly but that
272	# shouldn't be much of a problem (people shouldn't be using
273	# single quotes in filenames anyway :) )
274	params.append("%s='%s'" % (key, value))
275	return " ".join(params)
276
277	def execute_sql(self, dbname, user, sql):
278	"""Execute the sql statament
279
280	The user parameter us used as in connection_params. The dbname
281	parameter must be the name of a database in the cluster.
282	"""
283	conn = psycopg.connect("dbname=%s " % dbname
284	+ self.connection_string(user))
285	cursor = conn.cursor()
286	cursor.execute(sql)
287	conn.commit()
288	conn.close()
289
290	def require_authentication(self, required):
291	"""Switch authentication requirements on or off
292
293	When started for the first time no passwords are required. Some
294	tests want to explicitly test whether Thuban's password
295	infrastructure works and switch password authentication on
296	explicitly. When switching it on, there should be a
297	corresponding call to switch it off again in the test case'
298	tearDown method or in a finally: block.
299	"""
300	if required:
301	contents = "local all password\n"
302	else:
303	contents = "local all trust\n"
304	f = open(os.path.join(self.dbdir, "pg_hba.conf"), "w")
305	f.write(contents)
306	f.close()
307	run_command(["pg_ctl", "-D", self.dbdir, "reload"],
308	os.path.join(self.dbdir, "pg_ctl-reload.log"))
309
310
311	def create_user(self, username, password):
312	"""Create user username with password in the database"""
313	self.execute_sql("template1", "admin",
314	"CREATE USER %s PASSWORD '%s';" % (username,password))
315
316	def alter_user(self, username, password):
317	"""Change the user username's password in the database"""
318	self.execute_sql("template1", "admin",
319	"ALTER USER %s PASSWORD '%s';" % (username,password))
320
321
322	class PostGISDatabase:
323
324	"""A PostGIS database in a PostgreSQLServer"""
325
326	def __init__(self, server, postgis_sql, dbname, tables = None):
327	"""Initialize the PostGISDatabase
328
329	Parameters:
330
331	server -- The PostgreSQLServer instance containing the
332	database
333
334	postgis_sql -- Filename of the postgis.sql file with the
335	postgis initialization code
336
337	dbname -- The name of the database
338
339	tables -- Optional description of tables to create in the
340	new database. If given it should be a list of
341	(tablename, shapefilename) pairs meaning that a table
342	tablename will be created with the contents of the given
343	shapefile or (tablename, shapefilename, extraargs)
344	triples. The extraargs should be a list of key, value
345	pairs to use as keyword arguments to upload_shapefile.
346	"""
347	self.server = server
348	self.postgis_sql = postgis_sql
349	self.dbname = dbname
350	self.tables = tables
351
352	def initdb(self):
353	"""Remove the old db directory and create and initialize a new database
354	"""
355	run_command(["createdb", "-p", str(self.server.port),
356	"-h", self.server.host, "-U", self.server.admin_name,
357	self.dbname],
358	os.path.join(self.server.dbdir, "createdb.log"))
359	run_command(["createlang", "-p", str(self.server.port),
360	"-h", self.server.host, "-U", self.server.admin_name,
361	"plpgsql", self.dbname],
362	os.path.join(self.server.dbdir, "createlang.log"))
363	# for some reason psql doesn't exit with an error code if the
364	# file given as -f doesn't exist, so we check manually by trying
365	# to open it before we run psql
366	f = open(self.postgis_sql)
367	f.close()
368	del f
369	run_command(["psql", "-f", self.postgis_sql, "-d", self.dbname,
370	"-p", str(self.server.port), "-h", self.server.host,
371	"-U", self.server.admin_name],
372	os.path.join(self.server.dbdir, "psql.log"))
373
374	self.server.execute_sql(self.dbname, "admin",
375	"GRANT SELECT ON geometry_columns TO PUBLIC;")
376
377	if self.tables is not None:
378	def unpack(item):
379	extra = {"force_wkt_type": None, "gid_offset": 0}
380	if len(info) == 2:
381	tablename, shapefile = info
382	else:
383	tablename, shapefile, kw = info
384	for key, val in kw:
385	extra[key] = val
386	return tablename, shapefile, extra
387
388	for info in self.tables:
389	tablename, shapefile, kw = unpack(info)
390	upload_shapefile(shapefile, self, tablename, **kw)
391
392	def has_data(self, tables):
393	return self.tables == tables
394
395
396	def find_postgis_sql():
397	"""Return the name of the postgis_sql file
398
399	A postgis installation usually has the postgis_sql file in
400	PostgreSQL's datadir (i.e. the directory where PostgreSQL keeps
401	static files, not the directory containing the databases).
402	Unfortunately there's no way to determine the name of this directory
403	with pg_config so we assume here that it's
404	$bindir/../share/postgresql/.
405	"""
406	bindir = run_config_script("pg_config --bindir").strip()
407	return os.path.join(bindir, "..", "share", "postgresql",
408	"contrib", "postgis.sql")
409
410	_postgres_server = None
411	def get_test_server():
412	"""Return the test database server object.
413
414	If it doesn't exist yet, create it first.
415
416	The server will use the directory postgis under the temp dir (as
417	defined by support.create_temp_dir()) for the database cluster.
418	Sockets will be created in tempdir.
419	"""
420	global _postgres_server
421	if _postgres_server is None:
422	tempdir = support.create_temp_dir()
423	dbdir = os.path.join(tempdir, "postgis")
424	socket_dir = tempdir
425
426	_postgres_server = PostgreSQLServer(dbdir, 6543, find_postgis_sql(),
427	socket_dir = socket_dir)
428	_postgres_server.createdb()
429
430	return _postgres_server
431
432	def shutdown_test_server():
433	"""Shutdown the test server if it is running"""
434	global _postgres_server
435	if _postgres_server is not None:
436	_postgres_server.shutdown()
437	_postgres_server = None
438
439
440	def reason_for_not_running_tests():
441	"""
442	Determine whether postgis tests can be run and return a reason they can't
443
444	There's no fool-proof way to reliably determine this short of
445	actually running the tests but we try the following here:
446
447	- test whether pg_ctl --help can be run successfully
448	- test whether the postgis_sql can be opened
449	The name of the postgis_sql file is determined by find_postgis_sql()
450	- psycopg can be imported successfully.
451	"""
452	try:
453	run_command(["pg_ctl", "--help"], None)
454	except RuntimeError:
455	return "Can't run PostGIS tests because pg_ctl fails"
456
457	try:
458	postgis_sql = find_postgis_sql()
459	except:
460	return "Can't run PostGIS tests because postgis.sql can't be found"
461
462	try:
463	f = open(postgis_sql)
464	f.close()
465	except:
466	return "Can't run PostGIS tests because postgis.sql can't be opened"
467
468	# The test for psycopg was already done when this module was
469	# imported so we only have to check whether it was successful
470	if psycopg is None:
471	return "Can't run PostGIS tests because psycopg can't be imported"
472
473	return ""
474
475
476	_cannot_run_postgis_tests = None
477	def skip_if_no_postgis():
478	global _cannot_run_postgis_tests
479	if _cannot_run_postgis_tests is None:
480	_cannot_run_postgis_tests = reason_for_not_running_tests()
481	if _cannot_run_postgis_tests:
482	raise support.SkipTest(_cannot_run_postgis_tests)
483
484	def coords_to_point(coords):
485	"""Return string with a WKT representation of the point in coords"""
486	x, y = coords[0]
487	return "POINT(%r %r)" % (x, y)
488
489	def coords_to_polygon(coords):
490	"""Return string with a WKT representation of the polygon in coords"""
491	poly = []
492	for ring in coords:
493	poly.append(", ".join(["%r %r" % p for p in ring]))
494	return "POLYGON((%s))" % "), (".join(poly)
495
496	def coords_to_multilinestring(coords):
497	"""Return string with a WKT representation of the arc in coords"""
498	poly = []
499	for ring in coords:
500	poly.append(", ".join(["%r %r" % p for p in ring]))
501	return "MULTILINESTRING((%s))" % "), (".join(poly)
502
503	def coords_to_multipolygon(coords):
504	"""Return string with a WKT representation of the polygon in coords"""
505	poly = []
506	for ring in coords:
507	poly.append(", ".join(["%r %r" % p for p in ring]))
508	return "MULTIPOLYGON(((%s)))" % ")), ((".join(poly)
509
510	wkt_converter = {
511	"POINT": coords_to_point,
512	"MULTILINESTRING": coords_to_multilinestring,
513	"POLYGON": coords_to_polygon,
514	"MULTIPOLYGON": coords_to_multipolygon,
515	}
516
517	def upload_shapefile(filename, db, tablename, force_wkt_type = None,
518	gid_offset = 0):
519	"""Upload a shapefile into a new database table
520
521	Parameters:
522
523	filename -- The name of the shapefile
524
525	db -- The PostGISDatabase instance representing the database
526
527	tablename -- The name of the table to create and into which the data
528	is to be inserted
529
530	force_wkt_type -- If given the real WKT geometry type to use instead
531	of the default that would be chosen based on the type of
532	the shapefile
533
534	gid_offset -- A number to add to the shapeid to get the value for
535	the gid column (default 0)
536	"""
537	import dbflib, shapelib
538
539	# We build this map here because we need shapelib which can only be
540	# imported after support.initthuban has been called which we can't
541	# easily do in this module because it's imported by support.
542	shp_to_wkt = {
543	shapelib.SHPT_POINT: "POINT",
544	shapelib.SHPT_ARC: "MULTILINESTRING",
545	shapelib.SHPT_POLYGON: "POLYGON",
546	}
547
548	server = db.server
549	dbname = db.dbname
550	conn = psycopg.connect("dbname=%s " % dbname
551	+ db.server.connection_string("admin"))
552	cursor = conn.cursor()
553
554	shp = shapelib.ShapeFile(filename)
555	dbf = dbflib.DBFFile(filename)
556	typemap = {dbflib.FTString: "VARCHAR",
557	dbflib.FTInteger: "INTEGER",
558	dbflib.FTDouble: "DOUBLE PRECISION"}
559
560	insert_formats = ["%(gid)s"]
561	fields = ["gid INT"]
562	for i in range(dbf.field_count()):
563	ftype, name, width, prec = dbf.field_info(i)
564	fields.append("%s %s" % (name, typemap[ftype]))
565	insert_formats.append("%%(%s)s" % name)
566	stmt = "CREATE TABLE %s (\n %s\n);" % (tablename,
567	",\n ".join(fields))
568	cursor.execute(stmt)
569	#print stmt
570
571	numshapes, shapetype, mins, maxs = shp.info()
572	wkttype = shp_to_wkt[shapetype]
573	if force_wkt_type:
574	wkttype = force_wkt_type
575	convert = wkt_converter[wkttype]
576
577	cursor.execute("select AddGeometryColumn('%(dbname)s',"
578	"'%(tablename)s', 'the_geom', '-1', '%(wkttype)s', 2);"
579	% locals())
580
581	insert_formats.append("GeometryFromText(%(the_geom)s, -1)")
582
583	insert = ("INSERT INTO %s VALUES (%s)"
584	% (tablename, ", ".join(insert_formats)))
585
586	for i in range(numshapes):
587	data = dbf.read_record(i)
588	data["tablename"] = tablename
589	data["gid"] = i + gid_offset
590	data["the_geom"] = convert(shp.read_object(i).vertices())
591	#print insert % data
592	cursor.execute(insert, data)
593
594	cursor.execute("GRANT SELECT ON %s TO PUBLIC;" % tablename)
595
596	conn.commit()
Name	Value
svn:eol-style	native
svn:keywords	Author Date Id Revision