lsst.pipe.tasks ge1bc22dbe7+3fe987bc49
ingestPgsql.py
Go to the documentation of this file.
1import os
2
3from lsst.pex.config import ConfigurableField
4from lsst.pipe.tasks.ingest import IngestTask, IngestConfig, RegisterTask, RegistryContext, fakeContext
5from lsst.daf.persistence.registries import PgsqlRegistry
6
7try:
8 import psycopg2 as pgsql
9 havePgSql = True
10except ImportError:
11 havePgSql = False
12
13
15 """Context manager to provide a pgsql registry
16 """
17 def __init__(self, registryName, createTableFunc, forceCreateTables):
18 """Construct a context manager
19
20 @param registryName: Name of registry file
21 @param createTableFunc: Function to create tables
22 @param forceCreateTables: Force the (re-)creation of tables?
23 """
24 self.registryNameregistryName = registryName
25 data = PgsqlRegistry.readYaml(registryName)
26 self.connconnconn = pgsql.connect(host=data["host"], port=data["port"], user=data["user"],
27 password=data["password"], database=data["database"])
28 cur = self.connconnconn.cursor()
29
30 # Check for existence of tables
31 cur.execute("SELECT relname FROM pg_class WHERE relkind='r' AND relname='raw'")
32 rows = cur.fetchall()
33
34 if forceCreateTables or len(rows) == 0:
35 # Delete all tables and start over.
36 # Not simply doing "DROP SCHEMA" and "CREATE SCHEMA" because of permissions.
37 cur.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public'")
38 tables = cur.fetchall()
39 for tt in tables:
40 cur.execute("DROP TABLE %s CASCADE" % tt)
41 createTableFunc(self.connconnconn)
42
43 def __exit__(self, excType, excValue, traceback):
44 self.connconnconn.commit()
45 self.connconnconn.close()
46 return False # Don't suppress any exceptions
47
48
50 placeHolder = "%s"
51
52 def openRegistry(self, directory, create=False, dryrun=False):
53 """Open the registry and return the connection handle.
54
55 @param directory Directory in which the registry file will be placed
56 @param create Clobber any existing registry and create a new one?
57 @param dryrun Don't do anything permanent?
58 @return Database connection
59 """
60 if dryrun:
61 return fakeContext()
62 registryName = os.path.join(directory, "registry.pgsql")
63 return PgsqlRegistryContext(registryName, self.createTablecreateTablecreateTable, create)
64
65 def createTable(self, conn, table=None):
66 """Create the registry tables
67
68 One table (typically 'raw') contains information on all files, and the
69 other (typically 'raw_visit') contains information on all visits.
70
71 This method is required because there's a slightly different syntax
72 compared to SQLite (FLOAT instead of DOUBLE, SERIAL instead of
73 AUTOINCREMENT).
74
75 @param conn Database connection
76 @param table Name of table to create in database
77 """
78 if table is None:
79 table = self.config.table
80
81 typeMap = {'int': 'INT',
82 'double': 'FLOAT', # Defaults to double precision
83 }
84
85 cur = conn.cursor()
86 cmd = "CREATE TABLE %s (id SERIAL NOT NULL PRIMARY KEY, " % table
87 cmd += ",".join(["%s %s" % (col, typeMap.get(colType.lower(), 'text')) for
88 col, colType in self.config.columns.items()])
89 if len(self.config.unique) > 0:
90 cmd += ", UNIQUE(" + ",".join(self.config.unique) + ")"
91 cmd += ")"
92 cur.execute(cmd)
93
94 cmd = "CREATE TABLE %s_visit (" % self.config.table
95 cmd += ",".join(["%s %s" % (col, typeMap.get(self.config.columns[col].lower(), 'TEXT')) for
96 col in self.config.visit])
97 cmd += ", UNIQUE(" + ",".join(set(self.config.visit).intersection(set(self.config.unique))) + ")"
98 cmd += ")"
99 cur.execute(cmd)
100 del cur
101 conn.commit()
102
103
105 register = ConfigurableField(target=PgsqlRegisterTask, doc="Registry entry")
106
107
109 ConfigClass = PgsqlIngestConfig
def createTable(self, conn, table=None, forceCreateTables=False)
Definition: ingest.py:302
def createTable(self, conn, table=None)
Definition: ingestPgsql.py:65
def openRegistry(self, directory, create=False, dryrun=False)
Definition: ingestPgsql.py:52
def __init__(self, registryName, createTableFunc, forceCreateTables)
Definition: ingestPgsql.py:17
def __exit__(self, excType, excValue, traceback)
Definition: ingestPgsql.py:43