Python sqlalchemy.types.Enum() Examples

def test_should_enum_convert_enum():
    field = get_field(types.Enum(enum.Enum("TwoNumbers", ("one", "two"))))
    field_type = field.type()
    assert isinstance(field_type, graphene.Enum)
    assert field_type._meta.name == "TwoNumbers"
    assert hasattr(field_type, "ONE")
    assert not hasattr(field_type, "one")
    assert hasattr(field_type, "TWO")
    assert not hasattr(field_type, "two")

    field = get_field(types.Enum("one", "two", name="two_numbers"))
    field_type = field.type()
    assert isinstance(field_type, graphene.Enum)
    assert field_type._meta.name == "TwoNumbers"
    assert hasattr(field_type, "ONE")
    assert not hasattr(field_type, "one")
    assert hasattr(field_type, "TWO")
    assert not hasattr(field_type, "two") 

def test_enum_constraint_type_doesnt_double(self):
        m1 = MetaData()

        t1 = Table(
            "x",
            m1,
            Column("flag", Enum("a", "b", "c", create_constraint=True)),
        )
        eq_(
            len([c for c in t1.constraints if isinstance(c, CheckConstraint)]),
            1,
        )
        m2 = MetaData()
        t2 = t1.to_metadata(m2)

        eq_(
            len([c for c in t2.constraints if isinstance(c, CheckConstraint)]),
            1,
        ) 

def print_create_table(tables):
    app.config.from_object('config.default')
    database.init_app(app)

    engine = database.session.get_bind()

    for class_name in tables:
        cls = get_class(class_name)

        for c in cls.__table__.columns:
            if not isinstance(c.type, Enum):
                continue
            t = c.type
            sql = str(CreateEnumType(t).compile(engine))
            click.echo(sql.strip() + ';')

        for index in cls.__table__.indexes:
            sql = str(CreateIndex(index).compile(engine))
            click.echo(sql.strip() + ';')

        sql = str(CreateTable(cls.__table__).compile(engine))
        click.echo(sql.strip() + ';') 

def test_schematype_ck_name_enum(self):
        m1 = MetaData(
            naming_convention={"ck": "ck_%(table_name)s_%(constraint_name)s"}
        )

        u1 = Table(
            "user",
            m1,
            Column("x", Enum("a", "b", name="foo", create_constraint=True)),
        )
        eq_(
            [c for c in u1.constraints if isinstance(c, CheckConstraint)][
                0
            ].name,
            "foo",
        )
        # but is hit at compile time
        self.assert_compile(
            schema.CreateTable(u1),
            'CREATE TABLE "user" ('
            "x VARCHAR(1), "
            "CONSTRAINT ck_user_foo CHECK (x IN ('a', 'b'))"
            ")",
        ) 

def test_querying_table(metadata):
    """
    Create an object for test table.

    """

    # When using pytest-xdist, we don't want concurrent table creations
    # across test processes so we assign a unique name for table based on
    # the current worker id.
    worker_id = os.environ.get('PYTEST_XDIST_WORKER', 'master')
    return Table(
        'test_querying_table_' + worker_id, metadata,
        Column('id', types.Integer, autoincrement=True, primary_key=True),
        Column('serial', types.Integer, Sequence("serial_seq")),
        Column('t_string', types.String(60), onupdate='updated'),
        Column('t_list', types.ARRAY(types.String(60))),
        Column('t_enum', types.Enum(MyEnum)),
        Column('t_int_enum', types.Enum(MyIntEnum)),
        Column('t_datetime', types.DateTime()),
        Column('t_date', types.DateTime()),
        Column('t_interval', types.Interval()),
        Column('uniq_uuid', PG_UUID, nullable=False, unique=True, default=uuid4),
    ) 

def test_pep435_custom_sort_key(self):
        one, two, a_member, b_member = (
            self.one,
            self.two,
            self.a_member,
            self.b_member,
        )

        def sort_enum_key_value(value):
            return str(value.value)

        typ = Enum(self.SomeEnum, sort_key_function=sort_enum_key_value)
        is_(typ.sort_key_function, sort_enum_key_value)

        eq_(
            sorted([two, one, a_member, b_member], key=typ.sort_key_function),
            [one, two, a_member, b_member],
        ) 

def test_broken_enum_returns_blanks(self):
        t = Table(
            "enum_missing",
            self.metadata,
            Column("id", Integer, primary_key=True),
            Column("e1", sqltypes.Enum("one", "two", "three")),
            Column("e2", mysql.ENUM("one", "two", "three")),
        )
        t.create()

        with testing.db.connect() as conn:
            conn.execute(
                t.insert(), {"e1": "nonexistent", "e2": "nonexistent"}
            )
            conn.execute(t.insert(), {"e1": "", "e2": ""})
            conn.execute(t.insert(), {"e1": "two", "e2": "two"})
            conn.execute(t.insert(), {"e1": None, "e2": None})

            eq_(
                conn.execute(
                    select([t.c.e1, t.c.e2]).order_by(t.c.id)
                ).fetchall(),
                [("", ""), ("", ""), ("two", "two"), (None, None)],
            ) 

def test_enum_compile(self):
        e1 = Enum("x", "y", "z", name="somename")
        t1 = Table("sometable", MetaData(), Column("somecolumn", e1))
        self.assert_compile(
            schema.CreateTable(t1),
            "CREATE TABLE sometable (somecolumn " "ENUM('x','y','z'))",
        )
        t1 = Table(
            "sometable",
            MetaData(),
            Column(
                "somecolumn",
                Enum("x", "y", "z", native_enum=False, create_constraint=True),
            ),
        )
        self.assert_compile(
            schema.CreateTable(t1),
            "CREATE TABLE sometable (somecolumn "
            "VARCHAR(1), CHECK (somecolumn IN ('x', "
            "'y', 'z')))",
        ) 

def test_create_table(self, connection):
        metadata = self.metadata
        t1 = Table(
            "table",
            metadata,
            Column("id", Integer, primary_key=True),
            Column(
                "value", Enum("one", "two", "three", name="onetwothreetype")
            ),
        )
        t1.create(connection)
        t1.create(connection, checkfirst=True)  # check the create
        connection.execute(t1.insert(), value="two")
        connection.execute(t1.insert(), value="three")
        connection.execute(t1.insert(), value="three")
        eq_(
            connection.execute(t1.select().order_by(t1.c.id)).fetchall(),
            [(1, "two"), (2, "three"), (3, "three")],
        ) 

def test_array_plus_native_enum_create(self):
        m = MetaData()
        t = Table(
            "t",
            m,
            Column(
                "data_1",
                self.ARRAY(postgresql.ENUM("a", "b", "c", name="my_enum_1")),
            ),
            Column(
                "data_2",
                self.ARRAY(types.Enum("a", "b", "c", name="my_enum_2")),
            ),
        )

        t.create(testing.db)
        eq_(
            set(e["name"] for e in inspect(testing.db).get_enums()),
            set(["my_enum_1", "my_enum_2"]),
        )
        t.drop(testing.db)
        eq_(inspect(testing.db).get_enums(), []) 

def test_generate_multiple(self):
        """Test that the same enum twice only generates once
        for the create_all() call, without using checkfirst.

        A 'memo' collection held by the DDL runner
        now handles this.

        """
        metadata = self.metadata

        e1 = Enum("one", "two", "three", name="myenum")
        Table("e1", metadata, Column("c1", e1))

        Table("e2", metadata, Column("c1", e1))

        metadata.create_all(checkfirst=False)
        metadata.drop_all(checkfirst=False)
        assert "myenum" not in [
            e["name"] for e in inspect(testing.db).get_enums()
        ] 

def test_generate_alone_on_metadata(self):
        """Test that the same enum twice only generates once
        for the create_all() call, without using checkfirst.

        A 'memo' collection held by the DDL runner
        now handles this.

        """
        metadata = self.metadata

        Enum("one", "two", "three", name="myenum", metadata=self.metadata)

        metadata.create_all(checkfirst=False)
        assert "myenum" in [e["name"] for e in inspect(testing.db).get_enums()]
        metadata.drop_all(checkfirst=False)
        assert "myenum" not in [
            e["name"] for e in inspect(testing.db).get_enums()
        ] 

def test_generate_multiple_on_metadata(self):
        metadata = self.metadata

        e1 = Enum("one", "two", "three", name="myenum", metadata=metadata)

        t1 = Table("e1", metadata, Column("c1", e1))

        t2 = Table("e2", metadata, Column("c1", e1))

        metadata.create_all(checkfirst=False)
        assert "myenum" in [e["name"] for e in inspect(testing.db).get_enums()]
        metadata.drop_all(checkfirst=False)
        assert "myenum" not in [
            e["name"] for e in inspect(testing.db).get_enums()
        ]

        e1.create()  # creates ENUM
        t1.create()  # does not create ENUM
        t2.create()  # does not create ENUM 

def test_reflection(self):
        metadata = self.metadata
        etype = Enum(
            "four", "five", "six", name="fourfivesixtype", metadata=metadata
        )
        Table(
            "table",
            metadata,
            Column("id", Integer, primary_key=True),
            Column(
                "value", Enum("one", "two", "three", name="onetwothreetype")
            ),
            Column("value2", etype),
        )
        metadata.create_all()
        m2 = MetaData(testing.db)
        t2 = Table("table", m2, autoload=True)
        eq_(t2.c.value.type.enums, ["one", "two", "three"])
        eq_(t2.c.value.type.name, "onetwothreetype")
        eq_(t2.c.value2.type.enums, ["four", "five", "six"])
        eq_(t2.c.value2.type.name, "fourfivesixtype") 

def test_remain_on_table_metadata_wide(self):
        metadata = self.metadata

        e1 = Enum("one", "two", "three", name="myenum", metadata=metadata)
        table = Table("e1", metadata, Column("c1", e1))

        # need checkfirst here, otherwise enum will not be created
        assert_raises_message(
            sa.exc.ProgrammingError,
            '.*type "myenum" does not exist',
            table.create,
        )
        table.create(checkfirst=True)
        table.drop()
        table.create(checkfirst=True)
        table.drop()
        assert "myenum" in [e["name"] for e in inspect(testing.db).get_enums()]
        metadata.drop_all()
        assert "myenum" not in [
            e["name"] for e in inspect(testing.db).get_enums()
        ] 

def test_standalone_enum(self):
        metadata = MetaData(testing.db)
        etype = Enum(
            "four", "five", "six", name="fourfivesixtype", metadata=metadata
        )
        etype.create()
        try:
            assert testing.db.dialect.has_type(testing.db, "fourfivesixtype")
        finally:
            etype.drop()
            assert not testing.db.dialect.has_type(
                testing.db, "fourfivesixtype"
            )
        metadata.create_all()
        try:
            assert testing.db.dialect.has_type(testing.db, "fourfivesixtype")
        finally:
            metadata.drop_all()
            assert not testing.db.dialect.has_type(
                testing.db, "fourfivesixtype"
            ) 

def test_adapt_length(self):
        from sqlalchemy.dialects.postgresql import ENUM

        e1 = Enum("one", "two", "three", length=50, native_enum=False)
        eq_(e1.adapt(ENUM).length, 50)
        eq_(e1.adapt(Enum).length, 50)

        e1 = Enum("one", "two", "three")
        eq_(e1.length, 5)
        eq_(e1.adapt(ENUM).length, 5)
        eq_(e1.adapt(Enum).length, 5) 

def test_no_length_non_native(self):
        e = Enum("x", "y", "long", native_enum=False)
        eq_(e.length, len("long")) 

def test_non_native_constraint_custom_type(self):
        class Foob(object):
            def __init__(self, name):
                self.name = name

        class MyEnum(TypeDecorator):
            def __init__(self, values):
                self.impl = Enum(
                    *[v.name for v in values],
                    name="myenum",
                    native_enum=False,
                    create_constraint=True
                )

            # future method
            def process_literal_param(self, value, dialect):
                return value.name

            def process_bind_param(self, value, dialect):
                return value.name

        m = MetaData()
        t1 = Table("t", m, Column("x", MyEnum([Foob("a"), Foob("b")])))
        const = [c for c in t1.constraints if isinstance(c, CheckConstraint)][
            0
        ]

        self.assert_compile(
            AddConstraint(const),
            "ALTER TABLE t ADD CONSTRAINT myenum CHECK (x IN ('a', 'b'))",
            dialect="default",
        ) 

def test_create_metadata_bound_no_crash(self):
        m1 = self.metadata
        Enum("a", "b", "c", metadata=m1, name="ncenum")

        m1.create_all(testing.db) 

def test_array_of_enums(self, array_cls, enum_cls, connection):
        tbl = Table(
            "enum_table",
            self.metadata,
            Column("id", Integer, primary_key=True),
            Column(
                "enum_col",
                array_cls(enum_cls("foo", "bar", "baz", name="an_enum")),
            ),
        )

        if util.py3k:
            from enum import Enum

            class MyEnum(Enum):
                a = "aaa"
                b = "bbb"
                c = "ccc"

            tbl.append_column(
                Column("pyenum_col", array_cls(enum_cls(MyEnum)),),
            )

        self.metadata.create_all(connection)

        connection.execute(
            tbl.insert(), [{"enum_col": ["foo"]}, {"enum_col": ["foo", "bar"]}]
        )

        sel = select([tbl.c.enum_col]).order_by(tbl.c.id)
        eq_(
            connection.execute(sel).fetchall(), [(["foo"],), (["foo", "bar"],)]
        )

        if util.py3k:
            connection.execute(tbl.insert(), {"pyenum_col": [MyEnum.a]})
            sel = select([tbl.c.pyenum_col]).order_by(tbl.c.id.desc())
            eq_(connection.scalar(sel), [MyEnum.a]) 

def test_raises_non_native_enums(self, array_cls, enum_cls):
        Table(
            "my_table",
            self.metadata,
            Column(
                "my_col",
                array_cls(
                    enum_cls(
                        "foo",
                        "bar",
                        "baz",
                        name="my_enum",
                        create_constraint=True,
                        native_enum=False,
                    )
                ),
            ),
        )

        testing.assert_raises_message(
            CompileError,
            "PostgreSQL dialect cannot produce the CHECK constraint "
            "for ARRAY of non-native ENUM; please specify "
            "create_constraint=False on this Enum datatype.",
            self.metadata.create_all,
            testing.db,
        ) 

def test_inherit_schema_enum(self):
        m = MetaData()
        type_ = sqltypes.Enum("a", "b", "c", schema="q", inherit_schema=True)
        t1 = Table("x", m, Column("y", type_), schema="z")
        eq_(t1.c.y.type.schema, "z") 

def test_generic_w_some_other_variant(self, connection):
        some_table = Table(
            "some_table",
            self.metadata,
            Column(
                "data",
                Enum(
                    "one", "two", "three", name="my_enum", native_enum=True
                ).with_variant(Enum("four", "five", "six"), "mysql"),
            ),
        )

        assert "my_enum" not in [
            e["name"] for e in inspect(connection).get_enums()
        ]

        self.metadata.create_all(connection)

        assert "my_enum" in [
            e["name"] for e in inspect(connection).get_enums()
        ]

        connection.execute(some_table.insert(), {"data": "two"})

        self.metadata.drop_all(connection)

        assert "my_enum" not in [
            e["name"] for e in inspect(connection).get_enums()
        ] 

def test_generic_w_pg_variant(self, connection):
        some_table = Table(
            "some_table",
            self.metadata,
            Column(
                "data",
                Enum(
                    "one",
                    "two",
                    "three",
                    native_enum=True  # make sure this is True because
                    # it should *not* take effect due to
                    # the variant
                ).with_variant(
                    postgresql.ENUM("four", "five", "six", name="my_enum"),
                    "postgresql",
                ),
            ),
        )

        assert "my_enum" not in [
            e["name"] for e in inspect(connection).get_enums()
        ]

        self.metadata.create_all(connection)

        assert "my_enum" in [
            e["name"] for e in inspect(connection).get_enums()
        ]

        connection.execute(some_table.insert(), {"data": "five"})

        self.metadata.drop_all(connection)

        assert "my_enum" not in [
            e["name"] for e in inspect(connection).get_enums()
        ] 

def test_schema_reflection(self):
        metadata = self.metadata
        etype = Enum(
            "four",
            "five",
            "six",
            name="fourfivesixtype",
            schema="test_schema",
            metadata=metadata,
        )
        Table(
            "table",
            metadata,
            Column("id", Integer, primary_key=True),
            Column(
                "value",
                Enum(
                    "one",
                    "two",
                    "three",
                    name="onetwothreetype",
                    schema="test_schema",
                ),
            ),
            Column("value2", etype),
        )
        metadata.create_all()
        m2 = MetaData(testing.db)
        t2 = Table("table", m2, autoload=True)
        eq_(t2.c.value.type.enums, ["one", "two", "three"])
        eq_(t2.c.value.type.name, "onetwothreetype")
        eq_(t2.c.value2.type.enums, ["four", "five", "six"])
        eq_(t2.c.value2.type.name, "fourfivesixtype")
        eq_(t2.c.value2.type.schema, "test_schema") 

def test_schematype_ck_name_propagate_conv(self):
        m1 = MetaData(
            naming_convention={"ck": "ck_%(table_name)s_%(constraint_name)s"}
        )

        u1 = Table(
            "user",
            m1,
            Column(
                "x",
                Enum(
                    "a", "b", name=naming.conv("foo"), create_constraint=True
                ),
            ),
        )
        eq_(
            [c for c in u1.constraints if isinstance(c, CheckConstraint)][
                0
            ].name,
            "foo",
        )
        # but is hit at compile time
        self.assert_compile(
            schema.CreateTable(u1),
            'CREATE TABLE "user" ('
            "x VARCHAR(1), "
            "CONSTRAINT foo CHECK (x IN ('a', 'b'))"
            ")",
        ) 

def bind_processor(self, dialect):
        enum_proc = sqltypes.Enum.bind_processor(self, dialect)

        def process(value):
            raw_str = enum_proc(value)
            return raw_str
        return process 

def bind_processor(self, dialect):
        enum_proc = sqltypes.Enum.bind_processor(self, dialect)

        def process(value):
            raw_str = enum_proc(value)
            return raw_str
        return process 

def _on_metadata_create(self, target, bind, **kw):
            super(SchemaTypeTest.TrackEvents, self)._on_metadata_create(
                target, bind, **kw
            )
            self.evt_targets += (target,)

    # TODO: Enum and Boolean put TypeEngine first.  Changing that here
    # causes collection-mutate-while-iterated errors in the event system
    # since the hooks here call upon the adapted type.  Need to figure out
    # why Enum and Boolean don't have this problem.