Python sqlalchemy.select() Examples

def batch_load_by_email(
        cls, context, emails=None, *,
        domain_name=None,
        is_active=None,
        status=None,
    ) -> Sequence[Optional[User]]:
        async with context['dbpool'].acquire() as conn:
            query = (
                sa.select([users])
                .select_from(users)
                .where(users.c.email.in_(emails))
            )
            if domain_name is not None:
                query = query.where(users.c.domain_name == domain_name)
            if status is not None:
                query = query.where(users.c.status == UserStatus(status))
            elif is_active is not None:  # consider is_active field only if status is empty
                _statuses = ACTIVE_USER_STATUSES if is_active else INACTIVE_USER_STATUSES
                query = query.where(users.c.status.in_(_statuses))
            return await batch_result(
                context, conn, query, cls,
                emails, lambda row: row['email'],
            ) 

def check_credential(dbpool, domain: str, email: str, password: str) \
                          -> Any:
    async with dbpool.acquire() as conn:
        query = (sa.select([users])
                   .select_from(users)
                   .where((users.c.email == email) &
                          (users.c.domain_name == domain)))
        result = await conn.execute(query)
        row = await result.first()
        if row is None:
            return None
        if row['password'] is None:
            # user password is not set.
            return None
        try:
            if _verify_password(password, row['password']):
                return row
        except ValueError:
            return None
        return None 

def load_all_user(cls, context, access_key):
        async with context['dbpool'].acquire() as conn:
            query = (sa.select([keypairs.c.user_id])
                       .select_from(keypairs)
                       .where(keypairs.c.access_key == access_key))
            result = await conn.execute(query)
            row = await result.fetchone()
            user_id = row['user_id']
            j = sa.join(
                keypairs, keypair_resource_policies,
                keypairs.c.resource_policy == keypair_resource_policies.c.name
            )
            query = (sa.select([keypair_resource_policies])
                       .select_from(j)
                       .where((keypairs.c.user_id == user_id)))
            return [cls.from_row(context, r) async for r in conn.execute(query)] 

def batch_load_by_name_user(cls, context, names):
        async with context['dbpool'].acquire() as conn:
            access_key = context['access_key']
            j = sa.join(
                keypairs, keypair_resource_policies,
                keypairs.c.resource_policy == keypair_resource_policies.c.name
            )
            query = (sa.select([keypair_resource_policies])
                       .select_from(j)
                       .where((keypair_resource_policies.c.name.in_(names)) &
                              (keypairs.c.access_key == access_key))
                       .order_by(keypair_resource_policies.c.name))
            return await batch_result(
                context, conn, query, cls,
                names, lambda row: row['name'],
            ) 

def mutate(cls, root, info, name, props):
        data = {
            'name': name,
            'default_for_unspecified':
                DefaultForUnspecified[props.default_for_unspecified],
            'total_resource_slots': ResourceSlot.from_user_input(
                props.total_resource_slots, None),
            'max_concurrent_sessions': props.max_concurrent_sessions,
            'max_containers_per_session': props.max_containers_per_session,
            'idle_timeout': props.idle_timeout,
            'max_vfolder_count': props.max_vfolder_count,
            'max_vfolder_size': props.max_vfolder_size,
            'allowed_vfolder_hosts': props.allowed_vfolder_hosts,
        }
        insert_query = (keypair_resource_policies.insert().values(data))
        item_query = (
            keypair_resource_policies.select()
            .where(keypair_resource_policies.c.name == name))
        return await simple_db_mutate_returning_item(
            cls, info.context, insert_query,
            item_query=item_query, item_cls=KeyPairResourcePolicy) 

def load_count(
        context, *,
        scaling_group=None,
        status=None,
    ) -> int:
        async with context['dbpool'].acquire() as conn:
            query = (
                sa.select([sa.func.count(agents.c.id)])
                .select_from(agents)
                .as_scalar()
            )
            if scaling_group is not None:
                query = query.where(agents.c.scaling_group == scaling_group)
            if status is not None:
                status = AgentStatus[status]
                query = query.where(agents.c.status == status)
            result = await conn.execute(query)
            count = await result.fetchone()
            return count[0] 

def load_all(
        cls, context, *,
        scaling_group=None,
        status=None,
    ) -> Sequence[Agent]:
        async with context['dbpool'].acquire() as conn:
            query = (
                sa.select([agents])
                .select_from(agents)
            )
            if scaling_group is not None:
                query = query.where(agents.c.scaling_group == scaling_group)
            if status is not None:
                status = AgentStatus[status]
                query = query.where(agents.c.status == status)
            return [
                cls.from_row(context, row) async for row in conn.execute(query)
            ] 

def batch_load(
        cls, context, agent_ids, *,
        status=None,
    ) -> Sequence[Optional[Agent]]:
        async with context['dbpool'].acquire() as conn:
            query = (sa.select([agents])
                       .select_from(agents)
                       .where(agents.c.id.in_(agent_ids))
                       .order_by(agents.c.id))
            if status is not None:
                status = AgentStatus[status]
                query = query.where(agents.c.status == status)
            return await batch_result(
                context, conn, query, cls,
                agent_ids, lambda row: row['id'],
            ) 

def mutate(cls, root, info, name, props):
        if _rx_slug.search(name) is None:
            return cls(False, 'invalid name format. slug format required.', None)
        data = {
            'name': name,
            'description': props.description,
            'is_active': props.is_active,
            'total_resource_slots': ResourceSlot.from_user_input(
                props.total_resource_slots, None),
            'allowed_vfolder_hosts': props.allowed_vfolder_hosts,
            'allowed_docker_registries': props.allowed_docker_registries,
            'integration_id': props.integration_id,
        }
        insert_query = (
            domains.insert()
            .values(data)
        )
        item_query = domains.select().where(domains.c.name == name)
        return await simple_db_mutate_returning_item(
            cls, info.context, insert_query,
            item_query=item_query, item_cls=Domain) 

def load_count(cls, context, *,
                         domain_name=None, group_id=None, user_id=None):
        from .user import users
        async with context['dbpool'].acquire() as conn:
            j = sa.join(vfolders, users, vfolders.c.user == users.c.uuid)
            query = (
                sa.select([sa.func.count(vfolders.c.id)])
                .select_from(j)
                .as_scalar()
            )
            if domain_name is not None:
                query = query.where(users.c.domain_name == domain_name)
            if group_id is not None:
                query = query.where(vfolders.c.group == group_id)
            if user_id is not None:
                query = query.where(vfolders.c.user == user_id)
            result = await conn.execute(query)
            count = await result.fetchone()
            return count[0] 

def load_slice(cls, context, limit, offset, *,
                         domain_name=None, group_id=None, user_id=None,
                         order_key=None, order_asc=None):
        from .user import users
        async with context['dbpool'].acquire() as conn:
            if order_key is None:
                _ordering = vfolders.c.created_at
            else:
                _order_func = sa.asc if order_asc else sa.desc
                _ordering = _order_func(getattr(vfolders.c, order_key))
            j = sa.join(vfolders, users, vfolders.c.user == users.c.uuid)
            query = (
                sa.select([vfolders])
                .select_from(j)
                .order_by(_ordering)
                .limit(limit)
                .offset(offset)
            )
            if domain_name is not None:
                query = query.where(users.c.domain_name == domain_name)
            if group_id is not None:
                query = query.where(vfolders.c.group == group_id)
            if user_id is not None:
                query = query.where(vfolders.c.user == user_id)
            return [cls.from_row(context, r) async for r in conn.execute(query)] 

def batch_load_by_user(cls, context, user_uuids, *,
                                 domain_name=None, group_id=None):
        from .user import users
        async with context['dbpool'].acquire() as conn:
            # TODO: num_attached count group-by
            j = sa.join(vfolders, users, vfolders.c.user == users.c.uuid)
            query = (
                sa.select([vfolders])
                .select_from(j)
                .where(vfolders.c.user.in_(user_uuids))
                .order_by(sa.desc(vfolders.c.created_at))
            )
            if domain_name is not None:
                query = query.where(users.c.domain_name == domain_name)
            if group_id is not None:
                query = query.where(vfolders.c.group == group_id)
            return await batch_multiresult(
                context, conn, query, cls,
                user_uuids, lambda row: row['user']
            ) 

def load_count(cls, context, session_id, *,
                         role=None,
                         domain_name=None, group_id=None, access_key=None):
        async with context['dbpool'].acquire() as conn:
            query = (
                sa.select([sa.func.count(kernels.c.id)])
                .select_from(kernels)
                # TODO: use "owner session ID" when we implement multi-container session
                .where(kernels.c.id == session_id)
                .as_scalar()
            )
            if role is not None:
                query = query.where(kernels.c.role == role)
            if domain_name is not None:
                query = query.where(kernels.c.domain_name == domain_name)
            if group_id is not None:
                query = query.where(kernels.c.group_id == group_id)
            if access_key is not None:
                query = query.where(kernels.c.access_key == access_key)
            result = await conn.execute(query)
            count = await result.fetchone()
            return count[0] 

def batch_load_detail(cls, context, container_ids, *,
                                domain_name=None, access_key=None):
        async with context['dbpool'].acquire() as conn:
            j = (
                kernels
                .join(groups, groups.c.id == kernels.c.group_id)
                .join(users, users.c.uuid == kernels.c.user_uuid)
            )
            query = (
                sa.select([kernels, groups.c.name, users.c.email])
                .select_from(j)
                .where(
                    (kernels.c.id.in_(container_ids))
                ))
            if domain_name is not None:
                query = query.where(kernels.c.domain_name == domain_name)
            if access_key is not None:
                query = query.where(kernels.c.access_key == access_key)
            return await batch_result(
                context, conn, query, cls,
                container_ids, lambda row: row['id'],
            ) 

def load_count(cls, context, *,
                         domain_name=None, group_id=None, access_key=None,
                         status=None):
        if isinstance(status, str):
            status_list = [KernelStatus[s] for s in status.split(',')]
        elif isinstance(status, KernelStatus):
            status_list = [status]
        async with context['dbpool'].acquire() as conn:
            query = (
                sa.select([sa.func.count(kernels.c.id)])
                .select_from(kernels)
                .where(kernels.c.role == 'master')
                .as_scalar()
            )
            if domain_name is not None:
                query = query.where(kernels.c.domain_name == domain_name)
            if group_id is not None:
                query = query.where(kernels.c.group_id == group_id)
            if access_key is not None:
                query = query.where(kernels.c.access_key == access_key)
            if status is not None:
                query = query.where(kernels.c.status.in_(status_list))
            result = await conn.execute(query)
            count = await result.fetchone()
            return count[0] 

def batch_load_by_dependency(cls, context, session_ids):
        async with context['dbpool'].acquire() as conn:
            j = sa.join(
                kernels, kernel_dependencies,
                kernels.c.id == kernel_dependencies.c.depends_on,
            )
            query = (
                sa.select([kernels])
                .select_from(j)
                .where(
                    (kernels.c.role == 'master') &
                    (kernel_dependencies.c.kernel_id.in_(session_ids))
                )
            )
            return await batch_multiresult(
                context, conn, query, cls,
                session_ids, lambda row: row['id'],
            ) 

def batch_load_detail(cls, context, session_ids, *,
                                domain_name=None, access_key=None):
        async with context['dbpool'].acquire() as conn:
            j = (
                kernels
                .join(groups, groups.c.id == kernels.c.group_id)
                .join(users, users.c.uuid == kernels.c.user_uuid)
            )
            query = (
                sa.select([kernels, groups.c.name, users.c.email])
                .select_from(j)
                .where(
                    (kernels.c.role == 'master') &
                    (kernels.c.id.in_(session_ids))
                ))
            if domain_name is not None:
                query = query.where(kernels.c.domain_name == domain_name)
            if access_key is not None:
                query = query.where(kernels.c.access_key == access_key)
            return await batch_result(
                context, conn, query, cls,
                session_ids, lambda row: row['id'],
            ) 

def recalc_concurrency_used(db_conn: SAConnection, access_key: AccessKey) -> None:
    query = (
        sa.update(keypairs)
        .values(
            concurrency_used=(
                sa.select([sa.func.count(kernels.c.id)])
                .select_from(kernels)
                .where(
                    (kernels.c.access_key == access_key) &
                    (kernels.c.status.in_(USER_RESOURCE_OCCUPYING_KERNEL_STATUSES))
                )
                .as_scalar()
            ),
        )
        .where(keypairs.c.access_key == access_key)
    )
    await db_conn.execute(query) 

def load_all(
        cls, context, *,
        domain_name=None,
        is_active=None,
        limit=None,
    ) -> Sequence[KeyPair]:
        from .user import users
        async with context['dbpool'].acquire() as conn:
            j = sa.join(keypairs, users, keypairs.c.user == users.c.uuid)
            query = (
                sa.select([keypairs])
                .select_from(j)
            )
            if domain_name is not None:
                query = query.where(users.c.domain_name == domain_name)
            if is_active is not None:
                query = query.where(keypairs.c.is_active == is_active)
            if limit is not None:
                query = query.limit(limit)
            return [
                cls.from_row(context, row) async for row in conn.execute(query)
            ] 

def batch_load_by_email(
        cls, context, user_ids, *,
        domain_name=None, is_active=None,
    ) -> Sequence[Sequence[Optional[KeyPair]]]:
        from .user import users
        async with context['dbpool'].acquire() as conn:
            j = sa.join(keypairs, users, keypairs.c.user == users.c.uuid)
            query = (sa.select([keypairs])
                       .select_from(j)
                       .where(keypairs.c.user_id.in_(user_ids)))
            if domain_name is not None:
                query = query.where(users.c.domain_name == domain_name)
            if is_active is not None:
                query = query.where(keypairs.c.is_active == is_active)
            return await batch_multiresult(
                context, conn, query, cls,
                user_ids, lambda row: row['user_id'],
            ) 

def batch_load_by_ak(
        cls, context, access_keys, *,
        domain_name=None,
    ) -> Sequence[Optional[KeyPair]]:
        async with context['dbpool'].acquire() as conn:
            from .user import users
            j = sa.join(keypairs, users, keypairs.c.user == users.c.uuid)
            query = (
                sa.select([keypairs])
                .select_from(j)
                .where(
                    keypairs.c.access_key.in_(access_keys)
                )
            )
            if domain_name is not None:
                query = query.where(users.c.domain_name == domain_name)
            return await batch_result(
                context, conn, query, cls,
                access_keys, lambda row: row['access_key'],
            ) 

def _unreserve_agent_slots(
    db_conn: SAConnection,
    sess_ctx: PendingSession,
    agent_ctx: AgentAllocationContext,
) -> None:
    # Un-reserve agent slots, using a separate db txn.
    query = (
        sa.select([agents.c.occupied_slots], for_update=True)
        .select_from(agents)
        .where(agents.c.id == agent_ctx.agent_id))
    current_occupied_slots = await db_conn.scalar(query)
    query = (
        sa.update(agents)
        .values({
            'occupied_slots': current_occupied_slots - sess_ctx.requested_slots
        })
        .where(agents.c.id == agent_ctx.agent_id))
    await db_conn.execute(query) 

def _get_history_result_mapper(self, session, resource_type):
        mappers = self._resource_type_to_mappers(session, resource_type)
        resource_cls = mappers['resource']
        history_cls = mappers['history']

        resource_cols = {}
        history_cols = {}
        for col in sqlalchemy.inspect(history_cls).columns:
            history_cols[col.name] = col
            if col.name in ["revision", "revision_end"]:
                value = None if col.name == "revision_end" else -1
                resource_cols[col.name] = sqlalchemy.bindparam(
                    col.name, value, col.type).label(col.name)
            else:
                resource_cols[col.name] = getattr(resource_cls, col.name)
        s1 = sqlalchemy.select(history_cols.values())
        s2 = sqlalchemy.select(resource_cols.values())
        if resource_type != "generic":
            s1 = s1.where(history_cls.revision == ResourceHistory.revision)
            s2 = s2.where(resource_cls.id == Resource.id)
        union_stmt = sqlalchemy.union(s1, s2)
        stmt = union_stmt.alias("result")

        class Result(base.ResourceJsonifier, base.GnocchiBase):
            def __iter__(self):
                return iter((key, getattr(self, key)) for key in stmt.c.keys())

        sqlalchemy.orm.mapper(
            Result, stmt, primary_key=[stmt.c.id, stmt.c.revision],
            properties={
                'metrics': sqlalchemy.orm.relationship(
                    Metric,
                    primaryjoin=sqlalchemy.and_(
                        Metric.resource_id == stmt.c.id,
                        Metric.status == 'active'),
                    foreign_keys=Metric.resource_id)
            })

        return Result 

def upgrade():
    resource_type = op.create_table(
        'resource_type',
        sa.Column('name', sa.String(length=255), nullable=False),
        sa.PrimaryKeyConstraint('name'),
        mysql_charset='utf8',
        mysql_engine='InnoDB'
    )

    resource = sa.Table('resource', sa.MetaData(),
                        type_string_col("type", "resource"))
    op.execute(resource_type.insert().from_select(
        ['name'], sa.select([resource.c.type]).distinct()))

    for table in ["resource", "resource_history"]:
        op.alter_column(table, "type", new_column_name="old_type",
                        existing_type=type_enum)
        op.add_column(table, type_string_col("type", table))
        sa_table = sa.Table(table, sa.MetaData(),
                            type_string_col("type", table),
                            type_enum_col('old_type'))
        op.execute(sa_table.update().values(
            {sa_table.c.type: sa_table.c.old_type}))
        op.drop_column(table, "old_type")
        op.alter_column(table, "type", nullable=False,
                        existing_type=type_string) 

def to_characters_query(ctx):
    query = select([character_table.c.id])
    return [row.id for row in ctx[SA_ENGINE_KEY].execute(query)] 

def to_characters_query(ctx):
    query = select([character_table.c.id])
    async with ctx[SA_ENGINE_KEY].acquire() as conn:
        rows = await conn.execute(query)
        return [row.id async for row in rows] 

def to_actors_query(ctx):
    query = select([actor_table.c.id])
    async with ctx[SA_ENGINE_KEY].acquire() as conn:
        rows = await conn.execute(query)
        return [row.id async for row in rows] 

def to_characters_query(ctx):
    query = select([character_table.c.id])
    return [row.id for row in ctx[SA_ENGINE_KEY].execute(query)] 

def to_actors_query(ctx):
    query = select([actor_table.c.id])
    return [row.id for row in ctx[SA_ENGINE_KEY].execute(query)] 

def select_expr(self, ids):
        # TODO: make this optional, but enabled by default
        filtered_ids = [i for i in set(ids) if i is not None]
        if filtered_ids:
            return (
                sqlalchemy.select([self.from_column.label('from_column'),
                                   self.to_column.label('to_column')])
                .where(self.in_impl(self.from_column, filtered_ids))
            )
        else:
            return None