Python sqlalchemy.ARRAY Examples

def upgrade():
    agentstatus.create(op.get_bind())
    kernelstatus.create(op.get_bind())
    op.add_column('agents', sa.Column('lost_at', sa.DateTime(timezone=True), nullable=True))
    op.add_column('agents', sa.Column('status', sa.Enum('ALIVE', 'LOST', 'RESTARTING', 'TERMINATED', name='agentstatus'), nullable=False))
    op.create_index(op.f('ix_agents_status'), 'agents', ['status'], unique=False)
    op.add_column('kernels', sa.Column('agent_addr', sa.String(length=128), nullable=False))
    op.add_column('kernels', sa.Column('cpu_slot', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('gpu_slot', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('mem_slot', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('repl_in_port', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('repl_out_port', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('stdin_port', sa.Integer(), nullable=False))
    op.add_column('kernels', sa.Column('stdout_port', sa.Integer(), nullable=False))
    op.drop_column('kernels', 'allocated_cores')
    op.add_column('kernels', sa.Column('cpu_set', sa.ARRAY(sa.Integer), nullable=True))
    op.add_column('kernels', sa.Column('gpu_set', sa.ARRAY(sa.Integer), nullable=True))
    op.alter_column('kernels', column_name='status', type_=sa.Enum(*kernelstatus_choices, name='kernelstatus'),
                    postgresql_using='status::kernelstatus') 

def test_arrays_base(self, connection):
        metadata = self.metadata
        t1 = Table(
            "t",
            metadata,
            Column("x", sqltypes.ARRAY(Float)),
            Column("y", sqltypes.ARRAY(REAL)),
            Column("z", sqltypes.ARRAY(postgresql.DOUBLE_PRECISION)),
            Column("q", sqltypes.ARRAY(Numeric)),
        )
        metadata.create_all()
        connection.execute(
            t1.insert(), x=[5], y=[5], z=[6], q=[decimal.Decimal("6.4")]
        )
        row = connection.execute(t1.select()).first()
        eq_(row, ([5], [5], [6], [decimal.Decimal("6.4")])) 

def test_array_getitem_slice_type(self):
        m = MetaData()
        arrtable = Table(
            "arrtable",
            m,
            Column("intarr", postgresql.ARRAY(Integer)),
            Column("strarr", postgresql.ARRAY(String)),
        )

        # type affinity is Array...
        is_(arrtable.c.intarr[1:3].type._type_affinity, ARRAY)
        is_(arrtable.c.strarr[1:3].type._type_affinity, ARRAY)

        # but the slice returns the actual type
        assert isinstance(arrtable.c.intarr[1:3].type, postgresql.ARRAY)
        assert isinstance(arrtable.c.strarr[1:3].type, postgresql.ARRAY) 

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_array_literal_getitem_multidim(self):
        obj = postgresql.array(
            [postgresql.array([1, 2]), postgresql.array([3, 4])]
        )

        self.assert_compile(
            obj,
            "ARRAY[ARRAY[%(param_1)s, %(param_2)s], "
            "ARRAY[%(param_3)s, %(param_4)s]]",
        )
        self.assert_compile(
            obj[1],
            "(ARRAY[ARRAY[%(param_1)s, %(param_2)s], "
            "ARRAY[%(param_3)s, %(param_4)s]])[%(param_5)s]",
        )
        self.assert_compile(
            obj[1][0],
            "(ARRAY[ARRAY[%(param_1)s, %(param_2)s], "
            "ARRAY[%(param_3)s, %(param_4)s]])[%(param_5)s][%(param_6)s]",
        ) 

def test_arrays_pg(self, connection):
        metadata = self.metadata
        t1 = Table(
            "t",
            metadata,
            Column("x", postgresql.ARRAY(Float)),
            Column("y", postgresql.ARRAY(REAL)),
            Column("z", postgresql.ARRAY(postgresql.DOUBLE_PRECISION)),
            Column("q", postgresql.ARRAY(Numeric)),
        )
        metadata.create_all()
        connection.execute(
            t1.insert(), x=[5], y=[5], z=[6], q=[decimal.Decimal("6.4")]
        )
        row = connection.execute(t1.select()).first()
        eq_(row, ([5], [5], [6], [decimal.Decimal("6.4")])) 

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "github_installations",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("account_login", sa.String(), nullable=True),
        sa.Column("account_id", sa.Integer(), nullable=True),
        sa.Column("account_url", sa.String(), nullable=True),
        sa.Column("account_type", sa.String(), nullable=True),
        sa.Column("sender_id", sa.Integer(), nullable=True),
        sa.Column("sender_login", sa.String(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=True),
        sa.Column(
            "repositories",
            sa.ARRAY(sa.Integer(), as_tuple=sa.ForeignKey("git_projects.id")),
            nullable=True,
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    # ### end Alembic commands ### 

def downgrade():
    op.drop_column('kernels', 'stdout_port')
    op.drop_column('kernels', 'stdin_port')
    op.drop_column('kernels', 'repl_out_port')
    op.drop_column('kernels', 'repl_in_port')
    op.drop_column('kernels', 'mem_slot')
    op.drop_column('kernels', 'gpu_slot')
    op.drop_column('kernels', 'cpu_slot')
    op.drop_column('kernels', 'agent_addr')
    op.drop_index(op.f('ix_agents_status'), table_name='agents')
    op.drop_column('agents', 'status')
    op.drop_column('agents', 'lost_at')
    op.alter_column('kernels', column_name='status', type_=sa.String(length=64))
    op.add_column('kernels', sa.Column('allocated_cores', sa.ARRAY(sa.Integer), nullable=True))
    op.drop_column('kernels', 'cpu_set')
    op.drop_column('kernels', 'gpu_set')
    agentstatus.drop(op.get_bind())
    kernelstatus.drop(op.get_bind()) 

def test_percentile_cont_array(self):
        expr = func.percentile_cont(0.5, 0.7).within_group(
            column("data", Integer)
        )
        is_(expr.type._type_affinity, ARRAY)
        is_(expr.type.item_type._type_affinity, Integer) 

def test_where_has_all(self):
        self._test_where(
            self.hashcol.has_all(postgresql.array(["1", "2"])),
            "test_table.hash ?& ARRAY[%(param_1)s, %(param_2)s]",
        ) 

def test_percentile_cont_array_desc(self):
        expr = func.percentile_cont(0.5, 0.7).within_group(
            column("data", Integer).desc()
        )
        is_(expr.type._type_affinity, ARRAY)
        is_(expr.type.item_type._type_affinity, Integer) 

def test_array_functions_plus_getitem(self):
        """test parenthesizing of functions plus indexing, which seems
        to be required by PostgreSQL.

        """
        stmt = select(
            [
                func.array_cat(
                    array([1, 2, 3]),
                    array([4, 5, 6]),
                    type_=postgresql.ARRAY(Integer),
                )[2:5]
            ]
        )
        self.assert_compile(
            stmt,
            "SELECT (array_cat(ARRAY[%(param_1)s, %(param_2)s, %(param_3)s], "
            "ARRAY[%(param_4)s, %(param_5)s, %(param_6)s]))"
            "[%(param_7)s:%(param_8)s] AS anon_1",
        )

        self.assert_compile(
            func.array_cat(
                array([1, 2, 3]),
                array([4, 5, 6]),
                type_=postgresql.ARRAY(Integer),
            )[3],
            "(array_cat(ARRAY[%(param_1)s, %(param_2)s, %(param_3)s], "
            "ARRAY[%(param_4)s, %(param_5)s, %(param_6)s]))[%(array_cat_1)s]",
        ) 

def test_array_agg_generic(self):
        expr = func.array_agg(column("q", Integer))
        is_(expr.type.__class__, types.ARRAY)
        is_(expr.type.item_type.__class__, Integer) 

def define_tables(cls, metadata):
        class ProcValue(TypeDecorator):
            impl = cls.ARRAY(Integer, dimensions=2)

            def process_bind_param(self, value, dialect):
                if value is None:
                    return None
                return [[x + 5 for x in v] for v in value]

            def process_result_value(self, value, dialect):
                if value is None:
                    return None
                return [[x - 7 for x in v] for v in value]

        Table(
            "arrtable",
            metadata,
            Column("id", Integer, primary_key=True),
            Column("intarr", cls.ARRAY(Integer)),
            Column("strarr", cls.ARRAY(Unicode())),
            Column("dimarr", ProcValue),
        )

        Table(
            "dim_arrtable",
            metadata,
            Column("id", Integer, primary_key=True),
            Column("intarr", cls.ARRAY(Integer, dimensions=1)),
            Column("strarr", cls.ARRAY(Unicode(), dimensions=1)),
            Column("dimarr", ProcValue),
        ) 

def test_reflect_array_column(self):
        metadata2 = MetaData(testing.db)
        tbl = Table("arrtable", metadata2, autoload=True)
        assert isinstance(tbl.c.intarr.type, self.ARRAY)
        assert isinstance(tbl.c.strarr.type, self.ARRAY)
        assert isinstance(tbl.c.intarr.type.item_type, Integer)
        assert isinstance(tbl.c.strarr.type.item_type, String) 

def test_array_str_collation(self):
        m = self.metadata

        t = Table(
            "t",
            m,
            Column("data", sqltypes.ARRAY(String(50, collation="en_US"))),
        )

        t.create() 

def test_array_index_slice_exprs(self, connection):
        """test a variety of expressions that sometimes need parenthesizing"""

        stmt = select([array([1, 2, 3, 4])[2:3]])
        eq_(connection.execute(stmt).scalar(), [2, 3])

        stmt = select([array([1, 2, 3, 4])[2]])
        eq_(connection.execute(stmt).scalar(), 2)

        stmt = select([(array([1, 2]) + array([3, 4]))[2:3]])
        eq_(connection.execute(stmt).scalar(), [2, 3])

        stmt = select([array([1, 2]) + array([3, 4])[2:3]])
        eq_(connection.execute(stmt).scalar(), [1, 2, 4])

        stmt = select([array([1, 2])[2:3] + array([3, 4])])
        eq_(connection.execute(stmt).scalar(), [2, 3, 4])

        stmt = select(
            [
                func.array_cat(
                    array([1, 2, 3]),
                    array([4, 5, 6]),
                    type_=self.ARRAY(Integer),
                )[2:5]
            ]
        )
        eq_(connection.execute(stmt).scalar(), [2, 3, 4, 5]) 

def test_any_all_exprs_array(self, connection):
        stmt = select(
            [
                3
                == any_(
                    func.array_cat(
                        array([1, 2, 3]),
                        array([4, 5, 6]),
                        type_=self.ARRAY(Integer),
                    )
                )
            ]
        )
        eq_(connection.execute(stmt).scalar(), True) 

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_array_agg_array_datatype(self):
        expr = func.array_agg(column("data", ARRAY(Integer)))
        is_(expr.type._type_affinity, ARRAY)
        is_(expr.type.item_type._type_affinity, Integer) 

def test_where_has_any(self):
        self._test_where(
            self.hashcol.has_any(postgresql.array(["1", "2"])),
            "test_table.hash ?| ARRAY[%(param_1)s, %(param_2)s]",
        ) 

def _myarray_fixture(self):
        class MyArray(ARRAY):
            pass

        return MyArray 

def test_array_slice_index(self):
        col = column("x", postgresql.ARRAY(Integer))
        self.assert_compile(
            select([col[5:10]]),
            "SELECT x[%(x_1)s:%(x_2)s] AS anon_1",
            checkparams={"x_2": 10, "x_1": 5},
        ) 

def test_array_agg_array_literal_explicit_type(self):
        from sqlalchemy.dialects.postgresql import array

        expr = array([column("data", Integer), column("d2", Integer)])

        agg_expr = func.array_agg(expr, type_=ARRAY(Integer))
        is_(agg_expr.type._type_affinity, ARRAY)
        is_(agg_expr.type.item_type._type_affinity, Integer)

        self.assert_compile(
            agg_expr, "array_agg(ARRAY[data, d2])", dialect="postgresql"
        ) 

def test_array_index_map_dimensions(self):
        col = column("x", ARRAY(Integer, dimensions=3))
        is_(col[5].type._type_affinity, ARRAY)
        eq_(col[5].type.dimensions, 2)
        is_(col[5][6].type._type_affinity, ARRAY)
        eq_(col[5][6].type.dimensions, 1)
        is_(col[5][6][7].type._type_affinity, Integer) 

def test_array_getitem_single_type(self):
        m = MetaData()
        arrtable = Table(
            "arrtable",
            m,
            Column("intarr", ARRAY(Integer)),
            Column("strarr", ARRAY(String)),
        )
        is_(arrtable.c.intarr[1].type._type_affinity, Integer)
        is_(arrtable.c.strarr[1].type._type_affinity, String) 

def test_array_agg_array_literal_implicit_type(self):
        from sqlalchemy.dialects.postgresql import array, ARRAY as PG_ARRAY

        expr = array([column("data", Integer), column("d2", Integer)])

        assert isinstance(expr.type, PG_ARRAY)

        agg_expr = func.array_agg(expr)
        assert isinstance(agg_expr.type, PG_ARRAY)
        is_(agg_expr.type._type_affinity, ARRAY)
        is_(agg_expr.type.item_type._type_affinity, Integer)

        self.assert_compile(
            agg_expr, "array_agg(ARRAY[data, d2])", dialect="postgresql"
        ) 

def test_array_getitem_slice_type(self):
        m = MetaData()
        arrtable = Table(
            "arrtable",
            m,
            Column("intarr", ARRAY(Integer)),
            Column("strarr", ARRAY(String)),
        )
        is_(arrtable.c.intarr[1:3].type._type_affinity, ARRAY)
        is_(arrtable.c.strarr[1:3].type._type_affinity, ARRAY) 

def test_array_getitem_slice_type_dialect_level(self):
        MyArray = self._myarray_fixture()
        m = MetaData()
        arrtable = Table(
            "arrtable",
            m,
            Column("intarr", MyArray(Integer)),
            Column("strarr", MyArray(String)),
        )
        is_(arrtable.c.intarr[1:3].type._type_affinity, ARRAY)
        is_(arrtable.c.strarr[1:3].type._type_affinity, ARRAY)

        # but the slice returns the actual type
        assert isinstance(arrtable.c.intarr[1:3].type, MyArray)
        assert isinstance(arrtable.c.strarr[1:3].type, MyArray) 

def test_array_index_map_dimensions(self):
        col = column("x", postgresql.ARRAY(Integer, dimensions=3))
        is_(col[5].type._type_affinity, ARRAY)
        assert isinstance(col[5].type, postgresql.ARRAY)
        eq_(col[5].type.dimensions, 2)
        is_(col[5][6].type._type_affinity, ARRAY)
        assert isinstance(col[5][6].type, postgresql.ARRAY)
        eq_(col[5][6].type.dimensions, 1)
        is_(col[5][6][7].type._type_affinity, Integer)