lsst.pipe.tasks  13.0-66-gfbf2f2ce+5
ingestPgsql.py
Go to the documentation of this file.
1 from __future__ import absolute_import, division, print_function
2 
3 import os
4 
5 from lsst.pex.config import ConfigurableField
6 from lsst.pipe.tasks.ingest import IngestTask, IngestConfig, RegisterTask, RegistryContext, fakeContext
7 from lsst.daf.persistence.registries import PgsqlRegistry
8 
9 try:
10  import psycopg2 as pgsql
11  havePgSql = True
12 except ImportError:
13  havePgSql = False
14 
15 
17  """Context manager to provide a pgsql registry
18  """
19  def __init__(self, registryName, createTableFunc, forceCreateTables):
20  """Construct a context manager
21 
22  @param registryName: Name of registry file
23  @param createTableFunc: Function to create tables
24  @param forceCreateTables: Force the (re-)creation of tables?
25  """
26  self.registryName = registryName
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()
31 
32  # Check for existence of tables
33  cur.execute("SELECT relname FROM pg_class WHERE relkind='r' AND relname='raw'")
34  rows = cur.fetchall()
35 
36  if forceCreateTables or len(rows) == 0:
37  # Delete all tables and start over.
38  # Not simply doing "DROP SCHEMA" and "CREATE SCHEMA" because of permissions.
39  cur.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
40  tables = cur.fetchall()
41  for tt in tables:
42  cur.execute("DROP TABLE %s CASCADE" % tt)
43  createTableFunc(self.conn)
44 
45  def __exit__(self, excType, excValue, traceback):
46  self.conn.commit()
47  self.conn.close()
48  return False # Don't suppress any exceptions
49 
50 
52  placeHolder = "%s"
53 
54  def openRegistry(self, directory, create=False, dryrun=False):
55  """Open the registry and return the connection handle.
56 
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
61  """
62  if dryrun:
63  return fakeContext()
64  registryName = os.path.join(directory, "registry.pgsql")
65  return PgsqlRegistryContext(registryName, self.createTable, create)
66 
67  def createTable(self, conn, table=None):
68  """Create the registry tables
69 
70  One table (typically 'raw') contains information on all files, and the
71  other (typically 'raw_visit') contains information on all visits.
72 
73  This method is required because there's a slightly different syntax
74  compared to SQLite (FLOAT instead of DOUBLE, SERIAL instead of
75  AUTOINCREMENT).
76 
77  @param conn Database connection
78  @param table Name of table to create in database
79  """
80  if table is None:
81  table = self.config.table
82 
83  typeMap = {'int': 'INT',
84  'double': 'FLOAT', # Defaults to double precision
85  }
86 
87  cur = conn.cursor()
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) + ")"
93  cmd += ")"
94  cur.execute(cmd)
95 
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))) + ")"
100  cmd += ")"
101  cur.execute(cmd)
102  del cur
103  conn.commit()
104 
105 
107  register = ConfigurableField(target=PgsqlRegisterTask, doc="Registry entry")
108 
109 
111  ConfigClass = PgsqlIngestConfig
def createTable(self, conn, table=None)
Definition: ingest.py:284
def __init__(self, registryName, createTableFunc, forceCreateTables)
Definition: ingestPgsql.py:19
def openRegistry(self, directory, create=False, dryrun=False)
Definition: ingestPgsql.py:54
def __exit__(self, excType, excValue, traceback)
Definition: ingestPgsql.py:45
def createTable(self, conn, table=None)
Definition: ingestPgsql.py:67