1 from __future__
import absolute_import, division, print_function
5 from lsst.pex.config
import ConfigurableField
7 from lsst.daf.persistence.registries
import PgsqlRegistry
10 import psycopg2
as pgsql
17 """Context manager to provide a pgsql registry
19 def __init__(self, registryName, createTableFunc, forceCreateTables):
20 """Construct a context manager
22 @param registryName: Name of registry file
23 @param createTableFunc: Function to create tables
24 @param forceCreateTables: Force the (re-)creation of tables?
27 data = PgsqlRegistry.readYaml(registryName)
28 self.
conn = pgsql.connect(host=data[
"host"], port=data[
"port"], user=data[
"user"],
29 password=data[
"password"], database=data[
"database"])
30 cur = self.conn.cursor()
33 cur.execute(
"SELECT relname FROM pg_class WHERE relkind='r' AND relname='raw'")
36 if forceCreateTables
or len(rows) == 0:
39 cur.execute(
"SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
40 tables = cur.fetchall()
42 cur.execute(
"DROP TABLE %s CASCADE" % tt)
43 createTableFunc(self.
conn)
45 def __exit__(self, excType, excValue, traceback):
55 """Open the registry and return the connection handle.
57 @param directory Directory in which the registry file will be placed
58 @param create Clobber any existing registry and create a new one?
59 @param dryrun Don't do anything permanent?
60 @return Database connection
64 registryName = os.path.join(directory,
"registry.pgsql")
68 """Create the registry tables
70 One table (typically 'raw') contains information on all files, and the
71 other (typically 'raw_visit') contains information on all visits.
73 This method is required because there's a slightly different syntax
74 compared to SQLite (FLOAT instead of DOUBLE, SERIAL instead of
77 @param conn Database connection
78 @param table Name of table to create in database
81 table = self.config.table
83 typeMap = {
'int':
'INT',
88 cmd =
"CREATE TABLE %s (id SERIAL NOT NULL PRIMARY KEY, " % table
89 cmd +=
",".join([
"%s %s" % (col, typeMap.get(colType.lower(),
'text'))
for
90 col, colType
in self.config.columns.items()])
91 if len(self.config.unique) > 0:
92 cmd +=
", UNIQUE(" +
",".join(self.config.unique) +
")"
96 cmd =
"CREATE TABLE %s_visit (" % self.config.table
97 cmd +=
",".join([
"%s %s" % (col, typeMap.get(self.config.columns[col].lower(),
'TEXT'))
for
98 col
in self.config.visit])
99 cmd +=
", UNIQUE(" +
",".join(set(self.config.visit).intersection(set(self.config.unique))) +
")"
107 register = ConfigurableField(target=PgsqlRegisterTask, doc=
"Registry entry")
111 ConfigClass = PgsqlIngestConfig