Merge pull request #17 from portfoliome/uuid-ext

Uuid ext

Author: pmart123

postpy/_version.py

@@ -1,3 +1,3 @@
-version_info = (0, 0, 2)
+version_info = (0, 0, 3)
 
 __version__ = '.'.join(map(str, version_info))

postpy/admin.py

@@ -6,6 +6,7 @@
 
 from postpy.base import Table, Column, Database, PrimaryKey
 from postpy.ddl import compile_qualified_name
+from postpy.extensions import install_extension
 from postpy.sql import select_dict
 
 
@@ -69,12 +70,6 @@ def reflect_table(conn, table_name, schema='public'):
     return Table(table_name, columns, primary_key, schema=schema)
 
 
-def pg_stats_extension_statement():
-    """postgres stats extensions statement."""
-
-    return 'CREATE EXTENSION IF NOT EXISTS pg_stat_statements;'
-
-
 def reset(db_name):
     """Reset database."""
 
@@ -87,3 +82,20 @@ def reset(db_name):
         cursor.execute(db.create_statement())
     conn.close()
 
+
+def install_extensions(extensions, **connection_parameters):
+    """Install Postgres extension if available.
+
+    Notes
+    -----
+    - superuser is generally required for installing extensions.
+    - Currently does not support specific schema.
+    """
+
+    from postpy.connections import connect
+
+    conn = connect(**connection_parameters)
+    conn.autocommit = True
+
+    for extension in extensions:
+        install_extension(conn, extension)

postpy/base.py

@@ -1,7 +1,6 @@
 from collections import namedtuple
 
 from foil.formatters import format_repr
-from foil.filters import create_indexer
 
 from postpy.ddl import (
     compile_column, compile_qualified_name, compile_primary_key,
@@ -160,9 +159,8 @@ def order_table_columns(table: Table, column_names: list) -> Table:
     """Record table column(s) and primary key columns by specified order."""
 
     unordered_columns = table.column_names
-    index_order = [unordered_columns.index(name) for name in column_names]
-    indexer = create_indexer(index_order)
-    ordered_columns = indexer(table.columns)
+    index_order = (unordered_columns.index(name) for name in column_names)
+    ordered_columns = [table.columns[i] for i in index_order]
     ordered_pkey_names = [column for column in column_names
                           if column in table.primary_key_columns]
     primary_key = PrimaryKey(ordered_pkey_names)

postpy/extensions.py

@@ -0,0 +1,37 @@
+import psycopg2
+from psycopg2._psycopg import AsIs
+
+
+def install_extension(conn, extension: str):
+    """Install Postgres extension."""
+
+    query = 'CREATE EXTENSION IF NOT EXISTS %s;'
+
+    with conn.cursor() as cursor:
+        cursor.execute(query, (AsIs(extension),))
+
+    installed = check_extension(conn, extension)
+
+    if not installed:
+        raise psycopg2.ProgrammingError(
+            'Postgres extension failed installation.', extension
+        )
+
+
+def check_extension(conn, extension: str) -> bool:
+    """Check to see if an extension is installed."""
+
+    query = 'SELECT installed_version FROM pg_available_extensions WHERE name=%s;'
+
+    with conn.cursor() as cursor:
+        cursor.execute(query, (extension,))
+        result = cursor.fetchone()
+
+    if result is None:
+        raise psycopg2.ProgrammingError(
+            'Extension is not available for installation.', extension
+        )
+    else:
+        extension_version = result[0]
+
+        return True if extension_version else False

postpy/sql.py

@@ -78,11 +78,11 @@ def select_dict(conn, query: str, params=None, name=None, itersize=5000):
             yield result
 
 
-def select_each(conn, query: str, parameter_groups):
+def select_each(conn, query: str, parameter_groups, name=None):
     """Run select query for each parameter set in single transaction."""
 
     with conn:
-        with conn.cursor() as cursor:
+        with conn.cursor(name=name) as cursor:
             for parameters in parameter_groups:
                 cursor.execute(query, parameters)
                 yield cursor.fetchone()

postpy/uuids.py

@@ -0,0 +1,65 @@
+"""Configure psycopg2 to support UUID conversion."""
+
+import psycopg2.extras
+
+from postpy.admin import install_extensions
+
+
+CRYPTO_EXTENSION = 'pgcrypto'
+UUID_OSSP_EXTENSION = 'uuid-ossp'
+
+
+def register_client():
+    """Have psycopg2 marshall UUID objects automatically."""
+
+    psycopg2.extras.register_uuid()
+
+
+def register_crypto():
+    """Support for UUID's on server side.
+
+    Lighter dependency than uuid-ossp supporting
+    random_uuid_function for UUID generation.
+    """
+
+    install_extensions([CRYPTO_EXTENSION])
+
+
+def register_uuid():
+    """Support for UUID's on server side.
+
+    Notes
+    -----
+    uuid-ossp can be problematic on some platforms. See:
+    https://www.postgresql.org/docs/current/static/uuid-ossp.html
+    """
+
+    install_extensions([UUID_OSSP_EXTENSION])
+
+
+def random_uuid_function(schema=None):
+    """Cryptographic random UUID function.
+
+    Generates random database side UUID's.
+
+    Notes
+    -----
+    Lighter dependency than uuid-ossp, but higher
+    fragmentation on disk if used as auto-generating primary key UUID.
+    """
+
+    return '{}gen_random_uuid()'.format(_format_schema(schema))
+
+
+def uuid_sequence_function(schema=None):
+    """Sequential UUID generation.
+
+    Sequential UUID creation on database side offering
+    less table fragmentation issues when used as UUID primary key.
+    """
+
+    return '{}uuid_generate_v1mc()'.format(_format_schema(schema))
+
+
+def _format_schema(schema):
+    return '{}.'.format(schema) if schema else ''

tests/test_admin.py

@@ -1,7 +1,8 @@
 import unittest
 
 from postpy.admin import (get_user_tables, get_primary_keys,
-                          get_column_metadata, reflect_table, reset)
+                          get_column_metadata, install_extensions,
+                          reflect_table, reset)
 from postpy.base import Database, Column, PrimaryKey, Table
 from postpy.connections import connect
 from postpy.fixtures import PostgreSQLFixture
@@ -95,3 +96,21 @@ def tearDown(self):
             cursor.execute(self.db.drop_statement())
 
         self.conn.close()
+
+
+class TestExtensions(PostgreSQLFixture, unittest.TestCase):
+    @classmethod
+    def _prep(cls):
+        cls.pg_extension = 'sslinfo'
+        cls.conn.autocommit = True
+
+    def test_install_extensions(self):
+
+        install_extensions([self.pg_extension])
+
+    @classmethod
+    def _clean(cls):
+        statement = 'DROP EXTENSION IF EXISTS {};'.format(cls.pg_extension)
+
+        with cls.conn.cursor() as cursor:
+            cursor.execute(statement)

tests/test_base.py

@@ -230,13 +230,14 @@ def test_make_delete_table(self):
 
 class TestReOrderTableColumns(unittest.TestCase):
     def setUp(self):
+        self.maxDiff = None
         self.schema = 'foo'
         self.table_name = 'foobar'
         columns = table_columns()
         primary_key = table_primary_keys()
         self.table = Table(self.table_name, columns, primary_key, self.schema)
 
-    def reorder_columns(self):
+    def test_reorder_columns(self):
         column_names = ['state', 'city', 'population']
 
         expect_columns = [Column('state', 'CHAR(2)', False),

tests/test_extensions.py

@@ -0,0 +1,20 @@
+import unittest
+
+import psycopg2
+
+from postpy.fixtures import PostgreSQLFixture
+from postpy.extensions import check_extension
+
+
+class TestExtensions(PostgreSQLFixture, unittest.TestCase):
+    @classmethod
+    def _prep(cls):
+        cls.pg_extension = 'sslinfo'
+
+    def test_check_no_extension(self):
+        with self.assertRaises(psycopg2.ProgrammingError):
+            check_extension(self.conn, 'fake_extension')
+
+    def test_check_uninstalled_extension(self):
+
+        self.assertFalse(check_extension(self.conn, self.pg_extension))

tests/test_sql.py

@@ -99,3 +99,11 @@ def test_query_columns(self):
         result = list(sql.query_columns(self.conn, query))
 
         self.assertEqual(expected, result)
+        query = 'select * from generate_series(1,3) as col1 where col1=%s;'
+        parameter_groups = [(3,), (2,), (1,)]
+        Record = namedtuple('Record', 'col1')
+
+        expected = [Record(col1=3), Record(col1=2), Record(col1=1)]
+        result = list(sql.select_each(self.conn, query, parameter_groups))
+
+        self.assertEqual(expected, result)