def setupDatabase(database, user, password=None):
# hack on, use psycopg2 connection to register postgres label -> nsa label adaptation
import psycopg2
conn = psycopg2.connect(user=user, password=password, database=database)
cur = conn.cursor()
register_composite('label', cur, globally=True, factory=LabelComposite)
register_composite('security_attribute',
cur,
globally=True,
factory=SecuritAttributeComposite)
cur.execute("SELECT oid FROM pg_type WHERE typname = 'timestamptz';")
timestamptz_oid = cur.fetchone()[0]
DT = psycopg2.extensions.new_type((timestamptz_oid, ), "timestamptz",
castDatetime)
psycopg2.extensions.register_type(DT)
conn.close()
Registry.DBPOOL = adbapi.ConnectionPool('psycopg2',
user=user,
password=password,
database=database)
def get_pub_articles_word_freq(publisher):
conn = None
try:
params = config()
print('connecting to db...')
conn = psycopg2.connect(**params)
cur = conn.cursor()
register_composite('word_freq', cur)
cur.execute("""
SET TIME ZONE 'UTC';
SELECT
tokens.title,
tokens.publisher,
tokens.word_freq::word_freq[],
tokens.two_word_freq,
tokens.three_word_freq,
metadata.tags,
metadata.mod
FROM tokens INNER JOIN metadata ON tokens.title = metadata.title
WHERE tokens.publisher = '%s';
""" % (publisher,))
values = cur.fetchall()
cur.close()
index = []
for article in values:
# word_freq
wf = []
for value in article[2]:
cluster = {'word': value.word,
'frequency': str(value.frequency),
'relativity': str(value.relativity)}
wf.append(cluster)
article = {'title': article[0],
'publisher': article[1],
'mod': article[6].isoformat(),
'tags': article[5],
'word_freq': wf,
'2-word_freq': article[3],
'3-word_freq': article[4]}
index.append(article)
# write index to disk as to avoid to let `online-open` to crash the server
if (publisher == 'online-open'):
save_to_json.dump('online-open-wf', index)
else:
return index
except (Exception, psycopg2.DatabaseError) as error:
print('db error:', error)
finally:
if conn is not None:
conn.close()
print('db connection closed')
def connect(params=_PARAMS):
"""Connect to database using @params"""
conn = psycopg2.connect(**params)
with conn:
with conn.cursor() as cur:
cur.execute("SET search_path TO kbpo;")
register_composite('kbpo.score', conn, True)
return conn
def register_model(self, ModelSubclass, typname=None):
"""Register an ORM model.
:param ModelSubclass: the :py:class:`~postgres.orm.Model` subclass to
register with this :py:class:`~postgres.Postgres` instance
:param typname: a string indicating the Postgres type to register this
model for (``typname``, without an "e," is the name of the relevant
column in the underlying ``pg_type`` table). If :py:class:`None`,
we'll look for :py:attr:`ModelSubclass.typname`.
:raises: :py:exc:`~postgres.NotAModel`,
:py:exc:`~postgres.NoTypeSpecified`,
:py:exc:`~postgres.NoSuchType`,
:py:exc:`~postgres.AlreadyRegistered`
.. note::
See the :py:mod:`~postgres.orm` docs for instructions on
subclassing :py:class:`~postgres.orm.Model`.
"""
self._validate_model_subclass(ModelSubclass)
if typname is None:
typname = getattr(ModelSubclass, 'typname', None)
if typname is None:
raise NoTypeSpecified(ModelSubclass)
n = self.one( "SELECT count(*) FROM pg_type WHERE typname=%s"
, (typname,)
)
if n < 1:
# Could be more than one since we don't constrain by typnamespace.
# XXX What happens then?
raise NoSuchType(typname)
if typname in self.model_registry:
existing_model = self.model_registry[typname]
raise AlreadyRegistered(existing_model, typname)
self.model_registry[typname] = ModelSubclass
ModelSubclass.db = self
# register a composite
with self.get_connection() as conn:
cursor = conn.cursor()
name = typname
if sys.version_info[0] < 3:
name = name.encode('UTF-8')
register_composite( name
, cursor
, globally=True
, factory=self.DelegatingCaster
)
def register_model(self, ModelSubclass, typname=None):
"""Register an ORM model.
:param ModelSubclass: the :py:class:`~postgres.orm.Model` subclass to
register with this :py:class:`~postgres.Postgres` instance
:param typname: a string indicating the Postgres type to register this
model for (``typname``, without an "e," is the name of the relevant
column in the underlying ``pg_type`` table). If :py:class:`None`,
we'll look for :py:attr:`ModelSubclass.typname`.
:raises: :py:exc:`~postgres.NotAModel`,
:py:exc:`~postgres.NoTypeSpecified`,
:py:exc:`~postgres.NoSuchType`,
:py:exc:`~postgres.AlreadyRegistered`
.. note::
See the :py:mod:`~postgres.orm` docs for instructions on
subclassing :py:class:`~postgres.orm.Model`.
"""
self._validate_model_subclass(ModelSubclass)
if typname is None:
typname = getattr(ModelSubclass, 'typname', None)
if typname is None:
raise NoTypeSpecified(ModelSubclass)
n = self.one( "SELECT count(*) FROM pg_type WHERE typname=%s"
, (typname,)
)
if n < 1:
# Could be more than one since we don't constrain by typnamespace.
# XXX What happens then?
raise NoSuchType(typname)
if typname in self.model_registry:
existing_model = self.model_registry[typname]
raise AlreadyRegistered(existing_model, typname)
self.model_registry[typname] = ModelSubclass
ModelSubclass.db = self
# register a composite
with self.get_connection() as conn:
cursor = conn.cursor()
name = typname
if sys.version_info[0] < 3:
name = name.encode('UTF-8')
register_composite( name
, cursor
, globally=True
, factory=self.DelegatingCaster
)
def get_allarticles_word_freq():
conn = None
try:
params = config()
print('connecting to db...')
conn = psycopg2.connect(**params)
cur = conn.cursor()
register_composite('word_freq', cur)
#-- values
cur.execute("""
SET TIME ZONE 'UTC';
SELECT
tokens.title,
tokens.publisher,
tokens.word_freq::word_freq[],
tokens.two_word_freq,
tokens.three_word_freq,
metadata.tags,
metadata.mod
FROM tokens INNER JOIN metadata ON tokens.title = metadata.title;
""")
values = cur.fetchall()
cur.close()
index = []
for article in values:
# word_freq
wf = []
for value in article[2]:
cluster = {'word': value.word,
'frequency': str(value.frequency),
'relativity': str(value.relativity)}
wf.append(cluster)
article = {'title': article[0],
'publisher': article[1],
'mod': article[6].isoformat(),
'tags': article[5],
'word_freq': wf,
'2-word_freq': article[3],
'3-word_freq': article[4]}
index.append(article)
return index
except (Exception, psycopg2.DatabaseError) as error:
print('db error:', error)
finally:
if conn is not None:
conn.close()
print('db connection closed')
def register_composite(cls, connection):
"""
Register this CompositeType with Postgres.
If the CompositeType does not yet exist in the database, this will
fail. Hopefully a migration will come along shortly and create the
type in the database. If `retry` is True, this CompositeType will try
to register itself again after the type is created.
"""
with connection.temporary_connection() as cur:
register_composite(cls._meta.db_type, cur, globally=True)
def create_engine(self, app, uri, dbname=None):
self.uri = uri
self.dbname = dbname
conn = self.connect()
with app.open_resource('sql/schema.sql') as f:
conn.cursor().execute(f.read())
conn.commit()
register_adapter(dict, Json)
register_adapter(datetime, self._adapt_datetime)
register_composite('history', conn, globally=True)
from alerta.models.alert import History
register_adapter(History, HistoryAdapter)
def register_composite(cls, connection):
klass = cls()
db_type = klass.db_type(connection)
if db_type:
try:
cls.python_type = register_composite(
str(db_type),
connection.cursor().cursor,
globally=True,
factory=klass.factory_class()
).type
except psycopg2.ProgrammingError:
_missing_types[db_type] = cls
else:
def adapt_composite(composite):
# For safety, `composite_python_class` must have the same
# attributes as the namedtuple `python_type`'s fields, so
# that those can be escaped rather than relying on
# `__str__`.
return AsIs("(%s)::%s" % (
", ".join([
adapt(getattr(composite, field)).getquoted().decode('utf-8') for field in cls.python_type._fields
]), db_type
))
register_adapter(cls.composite_python_class, adapt_composite)
def register_composite(cls, connection):
klass = cls()
db_type = klass.db_type(connection)
if db_type:
try:
cls.python_type = register_composite(
str(db_type),
connection.cursor().cursor,
globally=True,
factory=klass.factory_class()).type
except psycopg2.ProgrammingError:
_missing_types[db_type] = cls
else:
def adapt_composite(composite):
# For safety, `composite_python_class` must have the same
# attributes as the namedtuple `python_type`'s fields, so
# that those can be escaped rather than relying on
# `__str__`.
return AsIs("(%s)::%s" % (", ".join([
adapt(getattr(composite,
field)).getquoted().decode('utf-8')
for field in cls.python_type._fields
]), db_type))
register_adapter(cls.composite_python_class, adapt_composite)
def setupDatabase(database, user, password=None):
# hack on, use psycopg2 connection to register postgres label -> nsa label adaptation
import psycopg2
conn = psycopg2.connect(user=user, password=password, database=database)
cur = conn.cursor()
register_composite('label', cur, globally=True, factory=LabelComposite)
cur.execute("SELECT oid FROM pg_type WHERE typname = 'timestamptz';")
timestamptz_oid = cur.fetchone()[0]
DT = psycopg2.extensions.new_type((timestamptz_oid,), "timestamptz", castDatetime)
psycopg2.extensions.register_type(DT)
conn.close()
Registry.DBPOOL = adbapi.ConnectionPool('psycopg2', user=user, password=password, database=database)
def check_type(self, repr_sql, fetch_type, known_types, expect):
"""
Checks composite type structure and format.
"""
cursor = connection.cursor()
if repr_sql:
for _type in known_types:
register_composite(str(_type), cursor.cursor, factory=TupleComposite)
sql = 'SELECT ROW{repr_sql}::{ftype}'.format(repr_sql=repr_sql, ftype=fetch_type)
cursor.execute(sql)
result = cursor.fetchone()[0]
self.assertEqual(result, expect)
else:
result = run_query("SELECT COUNT(*) FROM pg_type WHERE typname = %s",
[fetch_type])
self.assertEqual(result, [(0,)])
def create_engine(self, app, uri, dbname=None):
self.uri = uri
self.dbname = dbname
conn = self.connect()
with app.open_resource('sql/schema.sql') as f:
conn.cursor().execute(f.read())
conn.commit()
register_adapter(dict, Json)
register_adapter(datetime, self._adapt_datetime)
register_composite(
'history',
conn,
globally=True
)
from alerta.models.alert import History
register_adapter(History, adapt_history)
def register_composite(cls, connection):
"""
Register this CompositeType with Postgres.
If the CompositeType does not yet exist in the database, this will
fail. Hopefully a migration will come along shortly and create the
type in the database. If `retry` is True, this CompositeType will try
to register itself again after the type is created.
"""
LOGGER.debug("Registering composite type %s on connection %s",
cls.__name__, connection)
cls.registered_connection = connection
with connection.temporary_connection() as cur:
# This is what to do when the type is coming out of the database
register_composite(cls._meta.db_type, cur, globally=True,
factory=cls.Caster)
# This is what to do when the type is going in to the database
register_adapter(cls, QuotedCompositeType)
def test_non_dbapi_connection(self):
self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
conn = self.connect(connection_factory=RealDictConnection)
try:
register_composite("type_ii", conn)
curs = conn.cursor()
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()["x"], (1, 2))
finally:
conn.close()
conn = self.connect(connection_factory=RealDictConnection)
try:
curs = conn.cursor()
register_composite("type_ii", conn)
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()["x"], (1, 2))
finally:
conn.close()
def test_non_dbapi_connection(self):
from psycopg2.extras import RealDictConnection
from psycopg2.extras import register_composite
self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
conn = psycopg2.connect(dsn, connection_factory=RealDictConnection)
try:
register_composite('type_ii', conn)
curs = conn.cursor()
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()['x'], (1,2))
finally:
conn.close()
conn = psycopg2.connect(dsn, connection_factory=RealDictConnection)
try:
curs = conn.cursor()
register_composite('type_ii', conn)
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()['x'], (1,2))
finally:
conn.close()
def setupDatabase(database, user, password=None, host=None, connection_id_start=None):
# hack on, use psycopg2 connection to register postgres label -> nsa label adaptation
import psycopg2
conn = psycopg2.connect(user=user, password=password, database=database, host=host)
cur = conn.cursor()
register_composite('label', cur, globally=True, factory=LabelComposite)
register_composite('security_attribute', cur, globally=True, factory=SecuritAttributeComposite)
cur.execute("SELECT oid FROM pg_type WHERE typname = 'timestamptz';")
timestamptz_oid = cur.fetchone()[0]
DT = psycopg2.extensions.new_type((timestamptz_oid,), "timestamptz", castDatetime)
psycopg2.extensions.register_type(DT)
if connection_id_start:
r = cur.execute("INSERT INTO backend_connection_id (connection_id) VALUES (%s) ON CONFLICT DO NOTHING;", (connection_id_start,) )
conn.commit()
conn.close()
Registry.DBPOOL = adbapi.ConnectionPool('psycopg2', user=user, password=password, database=database, host=host)
def setupDatabase(database,
user,
password=None,
host=None,
connection_id_start=None):
# hack on, use psycopg2 connection to register postgres label -> nsa label adaptation
import psycopg2
conn = psycopg2.connect(user=user,
password=password,
database=database,
host=host)
cur = conn.cursor()
register_composite('label', cur, globally=True, factory=LabelComposite)
register_composite('security_attribute',
cur,
globally=True,
factory=SecuritAttributeComposite)
cur.execute("SELECT oid FROM pg_type WHERE typname = 'timestamptz';")
timestamptz_oid = cur.fetchone()[0]
DT = psycopg2.extensions.new_type((timestamptz_oid, ), "timestamptz",
castDatetime)
psycopg2.extensions.register_type(DT)
if connection_id_start:
r = cur.execute(
"INSERT INTO backend_connection_id (connection_id) VALUES (%s) ON CONFLICT DO NOTHING;",
(connection_id_start, ))
conn.commit()
conn.close()
Registry.DBPOOL = adbapi.ConnectionPool('psycopg2',
user=user,
password=password,
database=database,
host=host)
def test_non_dbapi_connection(self):
from psycopg2.extras import RealDictConnection
from psycopg2.extras import register_composite
self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
conn = psycopg2.connect(dsn, connection_factory=RealDictConnection)
try:
register_composite("type_ii", conn)
curs = conn.cursor()
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()["x"], (1, 2))
finally:
conn.close()
conn = psycopg2.connect(dsn, connection_factory=RealDictConnection)
try:
curs = conn.cursor()
register_composite("type_ii", conn)
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()["x"], (1, 2))
finally:
conn.close()
def register_composite(cls, connection, globally=True):
"""Register this composite type with psycopg2."""
# Sanity check: Are we using a dummy database?
# An application using django-pgfields may disable their database
# **entirely** for testing purposes. If this happens, this should
# be a no-op, rather than an error.
if connection.vendor in ('dummy', 'unknown'):
return
# Register the composite type with psycopg2.
return register_composite(str(cls.db_type()), connection.cursor(),
factory=cls.caster,
globally=globally,
)
def register_composite(cls, connection, globally=True):
"""Register this composite type with psycopg2."""
# Sanity check: Are we using a dummy database?
# An application using django-pgfields may disable their database
# **entirely** for testing purposes. If this happens, this should
# be a no-op, rather than an error.
if connection.vendor in ('dummy', 'unknown'):
return
# Register the composite type with psycopg2.
return register_composite(
str(cls.db_type()),
connection.cursor(),
factory=cls.caster,
globally=globally,
)
def register_composite(cls):
db_type = cls().db_type(connection)
if db_type:
try:
cls.python_type = register_composite(
db_type,
connection.cursor().cursor,
globally=True
).type
except ProgrammingError:
_missing_types[db_type] = cls
else:
def adapt_composite(composite):
return AsIs("(%s)::%s" % (
", ".join([
adapt(getattr(composite, field)).getquoted() for field in composite._fields
]), db_type
))
register_adapter(cls.python_type, adapt_composite)
def register_composite(cls):
db_type = cls.db_type
class Caster(CompositeCaster):
def make(self, values):
return cls(**dict(zip(self.attnames, values)))
if db_type in _registered_types:
raise ValueError(
'Type {} has already been registered.'.format(db_type))
try:
_registered_types[db_type] = register_composite(
db_type,
connection.cursor().cursor,
globally=True,
factory=Caster)
except (ProgrammingError, OperationalError) as exc:
_missing_types[db_type] = (cls, exc)
else:
register_adapter(cls, adapt_composite)
_missing_types.pop(db_type, None)
def test_subclass(self):
oid = self._create_type(
"type_isd",
[("anint", "integer"), ("astring", "text"), ("adate", "date")],
)
class DictComposite(CompositeCaster):
def make(self, values):
return dict(zip(self.attnames, values))
t = register_composite("type_isd", self.conn, factory=DictComposite)
self.assertEqual(t.name, "type_isd")
self.assertEqual(t.oid, oid)
curs = self.conn.cursor()
r = (10, "hello", date(2011, 1, 2))
curs.execute("select %s::type_isd;", (r, ))
v = curs.fetchone()[0]
self.assert_(isinstance(v, dict))
self.assertEqual(v["anint"], 10)
self.assertEqual(v["astring"], "hello")
self.assertEqual(v["adate"], date(2011, 1, 2))
def register_composite(cls, connection, globally=True):
"""Register this composite type with psycopg2."""
return register_composite(str(cls.db_type()), connection.cursor(),
factory=cls.caster,
globally=globally,
)
def register_type_globaly(cls):
cursor = connection.cursor()
register_composite(cls.type_name(), cursor, globally=True)
register_adapter(cls, lambda obj: AsIs(obj._as_sql()))
print "Registering %s type" % (cls.type_name())
from django.db import models
from django.db import connection
from psycopg2.extras import register_composite
from psycopg2.extensions import register_adapter, adapt, AsIs
Rgb = register_composite(
'rgb_color_value',
connection.cursor().cursor,
globally=True
).type
def rgb_adapter(value):
return AsIs("(%s, %s, %s)::rgb_color_value" % (
adapt(value.red).getquoted(),
adapt(value.green).getquoted(),
adapt(value.blue).getquoted()
))
register_adapter(Rgb, rgb_adapter)
class Rgb:
def __init__(self, red, green, blue):
self.red = red
self.green = green
self.blue = blue
class RgbField(models.Field):
def parse_rgb(self, value):
return Rgb(value.red, value.green, value.blue)
from django.contrib.postgres.fields import ArrayField
from django.urls import reverse
from mptt.models import MPTTModel, TreeForeignKey
from django.db.models import Sum
from contact.models import Customer
from django.contrib.contenttypes.fields import GenericForeignKey
from django.contrib.contenttypes.models import ContentType
from djmoney.models.fields import MoneyField
from moneyed import Money
from .utils.currency import Balance
from psycopg2.extras import register_composite
from psycopg2.extensions import register_adapter, adapt, AsIs
from . import managers
MoneyValue = register_composite('money_value',
connection.cursor().cursor,
globally=True).type
def moneyvalue_adapter(value):
return AsIs("(%s,%s)::money_value" %
(adapt(value.amount), adapt(value.currency.code)))
register_adapter(Money, moneyvalue_adapter)
class MoneyValueField(models.Field):
description = "wrapper for money_value composite type in postgres"
def from_db_value(self, value, expression, connection):