mtsql 1.11.14__py3-none-any.whl → 1.11.17__py3-none-any.whl

mt/sql/redshift/__init__.py

@@ -1,3 +1,6 @@
+from pkg_resources import DistributionNotFound, get_distribution, parse_version
+from sqlalchemy.dialects import registry  # noqa
+
 from .main import *
 
 __api__ = [
@@ -15,3 +18,26 @@ __api__ = [
     "drop_column",
     "conform",
 ]
+
+
+for package in ["psycopg2", "psycopg2-binary", "psycopg2cffi"]:
+    try:
+        if get_distribution(package).parsed_version < parse_version("2.5"):
+            raise ImportError("Minimum required version for psycopg2 is 2.5")
+        break
+    except DistributionNotFound:
+        pass
+
+registry.register("rs", "mt.sql.redshift.dialect", "RedshiftDialect_psycopg2")
+registry.register("rs.psycopg2", "mt.sql.redshift.dialect", "RedshiftDialect_psycopg2")
+registry.register(
+    "rs+psycopg2cffi",
+    "mt.sql.redshift.dialect",
+    "RedshiftDialect_psycopg2cffi",
+)
+
+registry.register(
+    "rs+redshift_connector",
+    "mt.sql.redshift.dialect",
+    "RedshiftDialect_redshift_connector",
+)

mt/sql/redshift/ddl.py

@@ -0,0 +1,264 @@
+import sqlalchemy as sa
+from sqlalchemy.ext import compiler as sa_compiler
+from sqlalchemy.schema import DDLElement
+
+
+def _check_if_key_exists(key):
+    return isinstance(key, sa.Column) or key
+
+
+def get_table_attributes(preparer,
+                         diststyle=None,
+                         distkey=None,
+                         sortkey=None,
+                         interleaved_sortkey=None):
+    """
+    Parse the table attributes into an acceptable string for Redshift,
+    checking for valid combinations of distribution options.
+
+    Parameters
+    ----------
+    preparer: RedshiftIdentifierPreparer, required
+        The preparer associated with the compiler, usually accessed through
+        compiler.preparer.  You can use a RedshiftDDLCompiler instance to
+        access it.
+    diststyle: str, optional
+        The diststle to use for the table attributes.  This must be one of:
+        ("ALL", "EVEN", "KEY").  If unset, Redshift passes AUTO. If KEY is used
+        a distkey argument must be provided.  Inversely, if a diststyle other
+        than KEY is provided, a distkey argument cannot be provided.
+    distkey: str or sqlalchemy.Column, optional
+        The distribution key to use the for the table attributes.  This can be
+        provided without any distsyle specified or with KEY diststyle
+        specified.
+    sortkey: str or sqlalchemy.Column (or iterable thereof), optional
+        The (compound) sort key(s) to use for the table attributes. Mutually
+        exclusive option from interleaved_sortkey.
+    interleaved_sortkey: str or sqlalchemy.Column (or iterable), optional
+        The (interleaved) sort key(s) to use for the table attributes. Mutually
+        exclusive option from sortkey.
+
+    Returns
+    -------
+    string
+        the table_attributes to append to a DDLElement, normally when creating
+        a table or materialized view.
+
+    Raises
+    ------
+    sqlalchemy.exc.ArgumentError
+        when an invalid diststyle is set,
+        when incompatable (diststyle, distkey) pairs are used,
+        when both sortkey and interleaved_sortkey is specified.
+    """
+    text = ""
+
+    has_distkey = _check_if_key_exists(distkey)
+    if diststyle:
+        diststyle = diststyle.upper()
+        if diststyle not in ('EVEN', 'KEY', 'ALL'):
+            raise sa.exc.ArgumentError(
+                u"diststyle {0} is invalid".format(diststyle)
+            )
+        if diststyle != 'KEY' and has_distkey:
+            raise sa.exc.ArgumentError(
+                u"DISTSTYLE EVEN/ALL is not compatible with a DISTKEY."
+            )
+        if diststyle == 'KEY' and not has_distkey:
+            raise sa.exc.ArgumentError(
+                u"DISTKEY specification is required for DISTSTYLE KEY"
+            )
+        text += " DISTSTYLE " + diststyle
+
+    if has_distkey:
+        if isinstance(distkey, sa.Column):
+            distkey = distkey.name
+        text += " DISTKEY ({0})".format(preparer.quote(distkey))
+
+    has_sortkey = _check_if_key_exists(sortkey)
+    has_interleaved = _check_if_key_exists(interleaved_sortkey)
+    if has_sortkey and has_interleaved:
+        raise sa.exc.ArgumentError(
+            "Parameters sortkey and interleaved_sortkey are "
+            "mutually exclusive; you may not specify both."
+        )
+
+    if has_sortkey or has_interleaved:
+        keys = sortkey if has_sortkey else interleaved_sortkey
+        if isinstance(keys, (str, sa.Column)):
+            keys = [keys]
+        keys = [key.name if isinstance(key, sa.Column) else key
+                for key in keys]
+        if has_interleaved:
+            text += " INTERLEAVED"
+        sortkey_string = ", ".join(preparer.quote(key)
+                                   for key in keys)
+        text += " SORTKEY ({0})".format(sortkey_string)
+    return text
+
+
+class CreateMaterializedView(DDLElement):
+    """
+    DDLElement to create a materialized view in Redshift where the distribution
+    options can be set.
+    SEE:
+    docs.aws.amazon.com/redshift/latest/dg/materialized-view-create-sql-command
+
+    This works for any selectable.  Consider the trivial example of this table:
+
+    >>> import sqlalchemy as sa
+    >>> from sqlalchemy_redshift.dialect import CreateMaterializedView
+    >>> engine = sa.create_engine('redshift+psycopg2://example')
+    >>> metadata = sa.MetaData()
+    >>> user = sa.Table(
+    ...     'user',
+    ...     metadata,
+    ...     sa.Column('id', sa.Integer, primary_key=True),
+    ...     sa.Column('name', sa.String)
+    ... )
+    >>> selectable = sa.select([user.c.id, user.c.name], from_obj=user)
+    >>> view = CreateMaterializedView(
+    ...     'materialized_view_of_users',
+    ...     selectable,
+    ...     distkey='id',
+    ...     sortkey='name'
+    ... )
+    >>> print(view.compile(engine))
+    <BLANKLINE>
+    CREATE MATERIALIZED VIEW materialized_view_of_users
+    DISTKEY (id) SORTKEY (name)
+    AS SELECT "user".id, "user".name
+    FROM "user"
+    <BLANKLINE>
+    <BLANKLINE>
+
+    The materialized view can take full advantage of Redshift's distributed
+    architecture via distribution styles and sort keys.
+
+    The CreateMaterializedView is a DDLElement, so it can be executed via any
+    execute() command, be it from an Engine, Session, or Connection.
+    """
+    def __init__(self, name, selectable, backup=True, diststyle=None,
+                 distkey=None, sortkey=None, interleaved_sortkey=None):
+        """
+        Parameters
+        ----------
+        name: str, required
+            the name of the materialized view to be created.
+        selectable: str, required
+            the sqlalchemy selectable to be the base query for the view.
+        diststyle: str, optional
+            The diststle to use for the table attributes.  This must be one of:
+            ("ALL", "EVEN", "KEY").  If unset, Redshift passes AUTO. If KEY is
+            used, a distkey argument must be provided.  Inversely, if
+            a diststyle other than KEY is provided, a distkey argument cannot
+            be provided.
+        distkey: str or sqlalchemy.Column, optional
+            The distribution key to use the for the table attributes.  This can
+            be provided without any distsyle specified or with KEY diststyle
+            specified.
+        sortkey: str or sqlalchemy.Column (or iterable thereof), optional
+            The (compound) sort key(s) to use for the table attributes.
+            Mutually exclusive option from interleaved_sortkey.
+        interleaved_sortkey: str or sqlalchemy.Column (or iterable), optional
+            The (interleaved) sort key(s) to use for the table attributes.
+            Mutually exclusive option from sortkey.
+        """
+        self.name = name
+        self.selectable = selectable
+        self.backup = backup
+        self.diststyle = diststyle
+        self.distkey = distkey
+        self.sortkey = sortkey
+        self.interleaved_sortkey = interleaved_sortkey
+
+
+@sa_compiler.compiles(CreateMaterializedView)
+def compile_create_materialized_view(element, compiler, **kw):
+    """
+    Returns the sql query that creates the materialized view
+    """
+
+    text = """\
+        CREATE MATERIALIZED VIEW {name}
+        {backup}
+        {table_attributes}
+        AS {selectable}\
+    """
+    table_attributes = get_table_attributes(
+        compiler.preparer,
+        diststyle=element.diststyle,
+        distkey=element.distkey,
+        sortkey=element.sortkey,
+        interleaved_sortkey=element.interleaved_sortkey
+    )
+    # Defaults to yes, so omit default cas3
+    backup = "" if element.backup else "BACKUP NO "
+    selectable = compiler.sql_compiler.process(element.selectable,
+                                               literal_binds=True)
+    text = text.format(
+        name=element.name,
+        backup=backup,
+        table_attributes=table_attributes,
+        selectable=selectable
+    )
+    # Clean it up to have no leading spaces
+    text = "\n".join([line.strip() for line in text.split("\n")
+                      if line.strip()])
+    return text
+
+
+class DropMaterializedView(DDLElement):
+    """
+    Drop the materialized view from the database.
+    SEE:
+    docs.aws.amazon.com/redshift/latest/dg/materialized-view-drop-sql-command
+
+    This undoes the create command, as expected:
+
+    >>> import sqlalchemy as sa
+    >>> from sqlalchemy_redshift.dialect import DropMaterializedView
+    >>> engine = sa.create_engine('redshift+psycopg2://example')
+    >>> drop = DropMaterializedView(
+    ...     'materialized_view_of_users',
+    ...     if_exists=True
+    ... )
+    >>> print(drop.compile(engine))
+    <BLANKLINE>
+    DROP MATERIALIZED VIEW IF EXISTS materialized_view_of_users
+    <BLANKLINE>
+    <BLANKLINE>
+
+    This can be included in any execute() statement.
+    """
+    def __init__(self, name, if_exists=False, cascade=False):
+        """
+        Build the DropMaterializedView DDLElement.
+
+        Parameters
+        ----------
+        name: str
+            name of the materialized view to drop
+        if_exists: bool, optional
+            if True, the IF EXISTS clause is added. This will make the query
+            successful even if the view does not exist, i.e. it lets you drop
+            a non-existant view. Defaults to False.
+        cascade: bool, optional
+            if True, the CASCADE clause is added. This will drop all
+            views/objects in the DB that depend on this materialized view.
+            Defaults to False.
+        """
+        self.name = name
+        self.if_exists = if_exists
+        self.cascade = cascade
+
+
+@sa_compiler.compiles(DropMaterializedView)
+def compile_drop_materialized_view(element, compiler, **kw):
+    """
+    Formats and returns the drop statement for materialized views.
+    """
+    text = "DROP MATERIALIZED VIEW {if_exists}{name}{cascade}"
+    if_exists = "IF EXISTS " if element.if_exists else ""
+    cascade = " CASCADE" if element.cascade else ""
+    return text.format(if_exists=if_exists, name=element.name, cascade=cascade)