Search code examples
python-3.xpostgresqlsqlalchemyflask-sqlalchemymarshmallow

Adding count of total rows through Marshmallow with @post_dump?


I need to add the quantity of rows returned in this query:

queryPostgres = db.text("""
            SELECT *, COUNT(*) OVER () as RowCount
            FROM (
                SELECT * ,
                    ( 3958.75 *
                    acos(sin(:lat1 / 57.2958) * sin( cast(latitude as double precision) / 57.2958) +
                         cos(:lat1 / 57.2958) * cos( cast(latitude as double precision) / 57.2958) *
                         cos( cast(longitude as double precision) / 57.2958 - :lon1/57.2958)))
                    as distanceInMiles
                FROM "job" ) zc
            WHERE zc.distanceInMiles < :dst
            ORDER BY zc.distanceInMiles
            LIMIT :per_page
            OFFSET :offset
        """)

        jobs = cls.query.\
            from_statement(queryPostgres). \
            params(lat1=float(lat1),
                   lon1=float(lon1),
                   dst=int(dst),
                   per_page=int(per_page),
                   offset=int(offset))
        return jobs

As you can see I added the RowCount column to have the total count of rows.

However as it is not part of my model, I wonder what should I do in Marshmallow so I could add the number of rows(in the RowCount column)?

I thought I could do it with Marshmallow @post_dump() , however I could not figure out how to do it .

For more clarity here is my schema.

class JobSchema(ma.ModelSchema):

    def validate_state(state):
        """Validate one of 55 USA states"""
        if state not in states:
            raise ValidationError(INVALID_US_STATE)

    def validate_zipCode(zip):
        if not zipcodes.is_real(zip):
            raise ValidationError(INVALID_ZIP_CODE)

    @pre_load
    def get_longitude_for_zipCode_and_TimeCreated(self, data):
        """ This method will pass valids long,lat and time_created
        values to each job created during a POST request"""
        # Getting zip from the request to obtain lat&lon from DB
        result = modelZipCode.getZipCodeDetails(data['zipCode'])
        print(result)
        if result is None:
            raise ValidationError(INVALID_ZIP_CODE_2)
        schema = ZipCodeSchema(exclude=('id'))
        zip, errors = schema.dump(result)
        if errors:
            raise ValidationError(INVALID_ZIP_CODE_3)
        else:
            data['longitude'] = zip['longitude']
            data['latitude'] = zip['latitude']
        data['time_created'] = str(datetime.datetime.utcnow())

    title = fields.Str(required=True, validate=[validate.Length(min=4, max=80)])
    city = fields.Str(required=True, validate=[validate.Length(min=5, max=100)])
    state = fields.Str(required=True, validate=validate_state)
    zipCode = fields.Str(required=True, validate=validate_zipCode)
    description = fields.Str(required=False, validate=[validate.Length(max=80)])
    narrative = fields.Str(required=False, validate=[validate.Length(max=250)])
    companyLogo = fields.Str(required=False, validate=[validate.Length(max=250)])
    companyName = fields.Str(required=True, validate=[validate.Length(min=5, max=250)])
    companyURL = fields.Str(required=True, validate=[validate.Length(min=4, max=100)])
    latitude = fields.Str(required=True)
    longitude = fields.Str(required=True)
    time_created = fields.DateTime()

    # We add a post_dump hook to add an envelope to responses
    @post_dump(pass_many=True)
    def wrap(self, data, many):
        #import pdb; pdb.set_trace()
        if len(data) >= 1:
           counter = data[0]['RowCount']

    return {
        data,
        counter
    }



    class Meta:
        model = modelJob

The most weird thing is that indeed my query is correctly returning the rowcount

enter image description here

Could some one please help me in finding out why I can not capture the rowcount key in the post_dump method ?


Solution

  • This need to be managed by the Marshmallow pre_dump or post_dump method.But indeed I decided to use an SQLAlchemy pagination methods as it gave me the total rows in the response.