path: root/generalresearch/managers/thl/survey.py

from collections import defaultdict
from datetime import datetime, timezone
from typing import Collection, List, Tuple, Optional

import pandas as pd
from more_itertools import chunked
from psycopg import sql

from generalresearch.managers.base import PostgresManager, Permission
from generalresearch.managers.thl.buyer import BuyerManager
from generalresearch.managers.thl.category import CategoryManager
from generalresearch.models import Source
from generalresearch.models.custom_types import SurveyKey
from generalresearch.models.thl.survey.model import (
    Survey,
    SurveyStat,
)
from generalresearch.pg_helper import PostgresConfig


class SurveyManager(PostgresManager):

    def __init__(
        self,
        pg_config: PostgresConfig,
        permissions: Collection[Permission] = None,
    ):
        super().__init__(pg_config=pg_config, permissions=permissions)
        self.buyer_manager = BuyerManager(pg_config=pg_config, permissions=permissions)
        self.category_manager = CategoryManager(pg_config=pg_config)

    def create_or_update(self, surveys: List[Survey]):
        """
        The only field that is checked for a possible update is `is_live`!
        """
        assert len({s.source for s in surveys}) == 1, "Only do one source at a time"
        source = surveys[0].source
        survey_ids = [s.survey_id for s in surveys]
        assert len(survey_ids) == len(set(survey_ids)), "duplicate survey_ids"

        # Handle the buyers
        buyer_codes = {s.buyer_code for s in surveys}
        self.buyer_manager.bulk_get_or_create(source=source, codes=buyer_codes)
        for s in surveys:
            s.buyer_id = self.buyer_manager.source_code_pk[s.buyer_natural_key]

        existing_surveys = self.filter_by_natural_key(
            source=source, survey_ids=survey_ids
        )
        existing_nks = {s.natural_key for s in existing_surveys}

        to_create = [
            survey for survey in surveys if survey.natural_key not in existing_nks
        ]
        if to_create:
            self.create_bulk(surveys=to_create)
            to_create_survey_ids = [s.survey_id for s in to_create]
            created_surveys = self.filter_by_natural_key(
                source=source, survey_ids=to_create_survey_ids
            )
            existing_surveys.extend(created_surveys)

        # Sometimes surveys get turned back on. Check that here
        potentially_update = [s for s in surveys if s.natural_key in existing_nks]
        existing_d = {s.survey_id: s for s in existing_surveys}
        to_update = []
        for s in potentially_update:
            if existing_d[s.survey_id].is_live != s.is_live:
                s.id = existing_d[s.survey_id].id
                to_update.append(s)
        if to_update:
            self.update_is_live(to_update)
        return {
            "survey_created_count": len(to_create),
            "survey_updated_count": len(to_update),
        }

    def create_bulk(self, surveys: List[Survey]):
        for chunk in chunked(surveys, 500):
            self.create_bulk_chunk(chunk)
        return None

    def create_bulk_chunk(self, surveys: List[Survey]):
        assert len(surveys) <= 500, "chunk me"

        query = """
        INSERT INTO marketplace_survey (
            source, survey_id, created_at, updated_at,
            is_live, is_recontact, buyer_id, eligibility_criteria
        ) VALUES (
            %(source)s, %(survey_id)s, %(created_at)s, %(updated_at)s,
            %(is_live)s, %(is_recontact)s, %(buyer_id)s, %(eligibility_criteria)s
        ) ON CONFLICT (source, survey_id) DO NOTHING;"""
        params = [s.model_dump_sql() for s in surveys]
        with self.pg_config.make_connection() as conn:
            with conn.cursor() as c:
                c.executemany(query=query, params_seq=params)
            conn.commit()
        return None

    def update_is_live(self, surveys: List[Survey]):
        ids_ON = [s.id for s in surveys if s.is_live]
        ids_OFF = [s.id for s in surveys if not s.is_live]
        query_ON = """
        UPDATE marketplace_survey
        SET is_live = TRUE, updated_at = NOW()
        WHERE id = ANY(%(ids)s);
        """
        query_OFF = """
        UPDATE marketplace_survey
        SET is_live = FALSE, updated_at = NOW()
        WHERE id = ANY(%(ids)s);
        """
        if ids_ON:
            self.pg_config.execute_write(
                query_ON,
                params={"ids": ids_ON},
            )
        if ids_OFF:
            self.pg_config.execute_write(
                query_OFF,
                params={"ids": ids_OFF},
            )

    def filter_by_keys(
        self,
        survey_keys: Collection[SurveyKey],
        include_categories: bool = False,
    ):
        assert len(survey_keys) <= 1000
        if len(survey_keys) == 0:
            return []

        params = dict()
        survey_source_ids = defaultdict(set)

        for sk in survey_keys:
            source, survey_id = sk.split(":")
            survey_source_ids[Source(source).value].add(survey_id)

        sk_filters = []
        for source, survey_ids in survey_source_ids.items():
            sk_filters.append(
                f"(s.source = '{source}' AND s.survey_id = ANY(%(survey_ids_{source})s))"
            )
            params[f"survey_ids_{source}"] = list(survey_ids)

        filter_str = f"WHERE ({' OR '.join(sk_filters)})"

        if include_categories:
            CATEGORY_JOIN = """
            LEFT JOIN LATERAL (
                SELECT
                    jsonb_agg(
                        jsonb_build_object(
                            'category',
                            jsonb_build_object(
                                'id', c.id,
                                'uuid', replace(c.uuid::text, '-', ''),
                                'label', c.label,
                                'path', c.path,
                                'adwords_vertical_id', c.adwords_vertical_id,
                                'parent_id', c.parent_id
                            ),
                            'strength', sc.strength
                        )
                        ORDER BY c.id
                    ) AS categories
                FROM marketplace_surveycategory sc
                JOIN marketplace_category c
                    ON c.id = sc.category_id
                WHERE sc.survey_id = s.id
            ) cat ON TRUE
            """

            query = f"""
            SELECT
                s.*,
                b.code as buyer_code,
                COALESCE(cat.categories, '[]'::jsonb) AS categories
            FROM marketplace_survey s
            LEFT JOIN marketplace_buyer b on s.buyer_id = b.id
            {CATEGORY_JOIN}
            {filter_str};
            """
        else:
            query = f"""
            SELECT s.*, b.code as buyer_code
            FROM marketplace_survey s
            LEFT JOIN marketplace_buyer b on s.buyer_id = b.id
            {filter_str};
            """

        res = self.pg_config.execute_sql_query(
            query,
            params=params,
        )
        return [Survey.model_validate(x) for x in res]

    def filter_by_natural_key(self, source: Source, survey_ids: Collection[str]):
        res = []
        for chunk in chunked(survey_ids, 1000):
            res.extend(self.filter_by_natural_key_chunk(source, chunk))
        return res

    def filter_by_natural_key_chunk(self, source: Source, survey_ids: Collection[str]):
        query = """
        SELECT id, source, survey_id, created_at, updated_at,
            is_live, is_recontact, buyer_id, eligibility_criteria
        FROM marketplace_survey
        WHERE source = %(source)s AND
            survey_id = ANY(%(survey_ids)s);
        """
        res = self.pg_config.execute_sql_query(
            query,
            params={"survey_ids": list(survey_ids), "source": source.value},
        )
        return [Survey.model_validate(x) for x in res]

    def filter_by_source_live(self, source: Source):
        """
        Return all live surveys for this source
        """
        query = """
        SELECT id, source, survey_id, created_at, updated_at,
            is_live, is_recontact, buyer_id, eligibility_criteria
        FROM marketplace_survey
        WHERE source = %(source)s AND is_live;
        """
        res = self.pg_config.execute_sql_query(query, params={"source": source.value})
        return [Survey.model_validate(x) for x in res]

    def filter_by_live(self, fields: Optional[List[str]] = None):
        """
        Return all live surveys
        """
        fields_default = """id, source, survey_id, created_at, updated_at,
            is_live, is_recontact, buyer_id, eligibility_criteria"""
        fields = ", ".join(fields) if fields else fields_default
        query = f"""
        SELECT {fields}
        FROM marketplace_survey
        WHERE is_live;
        """
        res = self.pg_config.execute_sql_query(query)
        return [Survey.model_validate(x) for x in res]

    def turn_off_by_natural_key(self, source: Source, survey_ids: Collection[str]):
        params = {"survey_ids": list(survey_ids), "source": source.value}
        query = """
        UPDATE marketplace_survey
        SET is_live = FALSE, updated_at = NOW()
        WHERE source = %(source)s AND
            survey_id = ANY(%(survey_ids)s)
        RETURNING id;
        """
        with self.pg_config.make_connection() as conn:
            with conn.cursor() as c:
                c.execute(query, params=params)
                survey_pks = [x["id"] for x in c.fetchall()]
            conn.commit()

        query = """
        UPDATE marketplace_surveystat
        SET survey_is_live = FALSE, updated_at = NOW()
        WHERE survey_is_live AND
            survey_id = ANY(%(survey_pks)s);
        """
        self.pg_config.execute_write(
            query,
            params={"survey_pks": survey_pks},
        )
        return None

    def update_surveys_categories(self, surveys: List[Survey] = None) -> None:
        for chunk in chunked(surveys, 500):
            self.update_surveys_categories_chunk(chunk)
        return None

    def update_surveys_categories_chunk(self, surveys: List[Survey] = None) -> None:
        assert len(surveys) <= 500, "chunk me"
        temp_table_sql = sql.SQL(
            """
        CREATE TEMP TABLE tmp_survey_categories (
            survey_id   bigint,
            category_id int,
            strength    float8
        ) ON COMMIT DROP;
        """
        )
        # noinspection SqlResolve
        insert_values_sql = sql.SQL(
            "INSERT INTO tmp_survey_categories VALUES (%s, %s, %s)"
        )
        # noinspection SqlResolve
        delete_sql = sql.SQL(
            """
        DELETE FROM marketplace_surveycategory sc
        WHERE NOT EXISTS (
            SELECT 1
            FROM tmp_survey_categories t
            WHERE t.survey_id = sc.survey_id
              AND t.category_id = sc.category_id
        )
        AND sc.survey_id IN (
            SELECT DISTINCT survey_id FROM tmp_survey_categories
        );"""
        )
        # noinspection SqlResolve
        upsert_sql = sql.SQL(
            """
        INSERT INTO marketplace_surveycategory (survey_id, category_id, strength)
        SELECT survey_id, category_id, strength
        FROM tmp_survey_categories
        ON CONFLICT (survey_id, category_id)
        DO UPDATE SET
            strength = EXCLUDED.strength;"""
        )

        rows = [
            (survey.id, c.category.id, c.strength)
            for survey in surveys
            for c in survey.categories
        ]
        with self.pg_config.make_connection() as conn:
            # noinspection PyArgumentList
            with conn.transaction():
                with conn.cursor() as c:
                    c.execute(temp_table_sql)
                    c.executemany(insert_values_sql, rows)
                    c.execute(delete_sql)
                    c.execute(upsert_sql)
            conn.commit()

    def get_survey_categories(self):
        query = """
        SELECT
            s.source, s.survey_id,
            jsonb_agg(
                jsonb_build_object(
                    'category_id', sc.category_id,
                    'strength', sc.strength
                )
            ) as categories
        FROM marketplace_survey s
        JOIN marketplace_surveycategory sc ON s.id = sc.survey_id
        WHERE is_live
        GROUP BY s.source, s.survey_id
        """
        return self.pg_config.execute_sql_query(query)


class SurveyStatManager(PostgresManager):
    KEYS = [
        "survey_id",
        "quota_id",
        "country_iso",
        "version",
        "cpi",
        "complete_too_fast_cutoff",
        "prescreen_conv_alpha",
        "prescreen_conv_beta",
        "conv_alpha",
        "conv_beta",
        "dropoff_alpha",
        "dropoff_beta",
        "completion_time_mu",
        "completion_time_sigma",
        "mobile_eligible_alpha",
        "mobile_eligible_beta",
        "desktop_eligible_alpha",
        "desktop_eligible_beta",
        "tablet_eligible_alpha",
        "tablet_eligible_beta",
        "long_fail_rate",
        "user_report_coeff",
        "recon_likelihood",
        "score_x0",
        "score_x1",
        "score",
        "updated_at",
        "survey_is_live",
        "survey_survey_id",
        "survey_source",
    ]

    SURVEY_STATS_COL_MAP = {
        "PRESCREEN_CONVERSION.alpha": "prescreen_conv_alpha",
        "PRESCREEN_CONVERSION.beta": "prescreen_conv_beta",
        "CONVERSION.alpha": "conv_alpha",
        "CONVERSION.beta": "conv_beta",
        "COMPLETION_TIME.mu": "completion_time_mu",
        "COMPLETION_TIME.sigma": "completion_time_sigma",
        "LONG_FAIL.value": "long_fail_rate",
        "USER_REPORT_COEFF.value": "user_report_coeff",
        "RECON_LIKELIHOOD.value": "recon_likelihood",
        "DROPOFF_RATE.alpha": "dropoff_alpha",
        "DROPOFF_RATE.beta": "dropoff_beta",
        "IS_MOBILE_ELIGIBLE.alpha": "mobile_eligible_alpha",
        "IS_MOBILE_ELIGIBLE.beta": "mobile_eligible_beta",
        "IS_DESKTOP_ELIGIBLE.alpha": "desktop_eligible_alpha",
        "IS_DESKTOP_ELIGIBLE.beta": "desktop_eligible_beta",
        "IS_TABLET_ELIGIBLE.alpha": "tablet_eligible_alpha",
        "IS_TABLET_ELIGIBLE.beta": "tablet_eligible_beta",
        "cpi": "cpi",
    }

    def __init__(
        self,
        pg_config: PostgresConfig,
        permissions: Collection[Permission] = None,
    ):
        super().__init__(pg_config=pg_config, permissions=permissions)
        self.survey_manager = SurveyManager(
            pg_config=pg_config, permissions=permissions
        )
        # self.ensure_surveystat_key_type()

    #
    # def ensure_surveystat_key_type(self):
    #     SQL = """
    #     DO $$
    #     BEGIN
    #         IF NOT EXISTS (
    #             SELECT 1
    #             FROM pg_type t
    #             JOIN pg_namespace n ON n.oid = t.typnamespace
    #             WHERE t.typname = 'surveystat_key'
    #               AND n.nspname = 'public'
    #         ) THEN
    #             CREATE TYPE public.surveystat_key AS (
    #                 survey_id   bigint,
    #                 quota_id    varchar(32),
    #                 country_iso varchar(2),
    #                 version     integer
    #             );
    #         END IF;
    #     END
    #     $$;"""
    #     with self.pg_config.make_connection() as conn:
    #         with conn.cursor() as c:
    #             c.execute(SQL)
    #         conn.commit()
    #     return None

    # def register_surveystat_key(self, conn):
    #     info = CompositeInfo.fetch(conn, "surveystat_key")
    #     info.register(conn)

    def update_or_create(self, survey_stats: List[SurveyStat]):
        """
        This manager is NOT responsible for creating surveys or buyers.
        It will check to make sure they exist
        """
        if len(survey_stats) == 0:
            return []
        assert all(s.survey_survey_id is not None for s in survey_stats)
        assert all(s.survey_source is not None for s in survey_stats)
        assert (
            len({s.survey_source for s in survey_stats}) == 1
        ), "Only do one source at a time"
        source = survey_stats[0].survey_source
        nks = [s.natural_key for s in survey_stats]
        assert len(nks) == len(set(nks)), "duplicate natural_keys"

        # Look up survey pks
        survey_ids = [s.survey_survey_id for s in survey_stats]
        surveys = self.survey_manager.filter_by_natural_key(
            source=source, survey_ids=survey_ids
        )
        nk_to_pk = {s.natural_key: s.id for s in surveys}
        for ss in survey_stats:
            try:
                ss.survey_id = nk_to_pk[ss.survey_natural_key]
            except KeyError as e:
                raise ValueError(
                    f"Survey {e.args[0]} does not exist. Must create surveys first"
                )
        # print(f"----aa-----: {datetime.now().isoformat()}")
        self.upsert_sql(survey_stats=survey_stats)
        # print(f"----ab-----: {datetime.now().isoformat()}")
        return None
        # keys = [s.unique_key for s in survey_stats]
        # print(keys[:4])
        # survey_stats = self.filter_by_unique_keys(keys)
        # print(len(survey_stats))
        # print(f"----ac-----: {datetime.now().isoformat()}")
        # # For testing/deterministic
        # survey_stats = sorted(survey_stats, key=lambda s: s.natural_key)
        # return survey_stats

    def upsert_sql(self, survey_stats: List[SurveyStat]):
        for chunk in chunked(survey_stats, 1000):
            self.upsert_sql_chunk(survey_stats=chunk)
        return None

    # def insert_sql(self, survey_stats: List[SurveyStat]):
    #     for chunk in chunked(survey_stats, 1000):
    #         self.insert_sql_chunk(survey_stats=chunk)
    #     return None
    #
    # def insert_sql_chunk(self, survey_stats: List[SurveyStat]):
    #     assert len(survey_stats) <= 1000, "chunk me"
    #     keys = self.keys
    #     keys_str = ", ".join(keys)
    #     values_str = ", ".join([f"%({k})s" for k in keys])
    #     unique_cols = ["survey_id", "quota_id", "country_iso", "version"]
    #     unique_cols_str = ", ".join(unique_cols)
    #
    #     query = f"""
    #     INSERT INTO marketplace_surveystat ({keys_str})
    #     VALUES ({values_str})
    #     ON CONFLICT ({unique_cols_str})
    #     DO NOTHING ;"""
    #     params = [ss.model_dump_sql() for ss in survey_stats]
    #     with self.pg_config.make_connection() as conn:
    #         with conn.cursor() as c:
    #             c.executemany(query=query, params_seq=params)
    #         conn.commit()
    #     return None

    def upsert_sql_chunk(self, survey_stats: List[SurveyStat]):
        assert len(survey_stats) <= 1000, "chunk me"
        keys = self.KEYS
        keys_str = ", ".join(keys)
        values_str = ", ".join([f"%({k})s" for k in keys])
        unique_cols = ["survey_id", "quota_id", "country_iso", "version"]
        unique_cols_str = ", ".join(unique_cols)
        update_cols = set(keys) - set(unique_cols) - {"updated_at", "is_live"}
        update_str = ", ".join(
            [f"{k} = EXCLUDED.{k}" for k in update_cols] + ["updated_at = NOW()"]
        )

        query = f"""
        INSERT INTO marketplace_surveystat ({keys_str})
        VALUES ({values_str})
        ON CONFLICT ({unique_cols_str})
        DO UPDATE SET {update_str};"""
        now = datetime.now(tz=timezone.utc)
        params = [ss.model_dump_sql() | {"updated_at": now} for ss in survey_stats]
        with self.pg_config.make_connection() as conn:
            with conn.cursor() as c:
                c.executemany(query=query, params_seq=params)
            conn.commit()
        return None

    def filter_by_unique_keys(self, keys: Collection[Tuple]):
        res = []
        for chunk in chunked(keys, 5000):
            res.extend(self.filter_by_unique_keys_chunk(chunk))
        return res

    def filter_by_unique_keys_chunk(self, keys: Collection[Tuple]):
        values_sql = ", ".join(["(%s, %s, %s, %s)"] * len(keys))
        query = f"""
        SELECT
            ss.*
        FROM marketplace_surveystat ss
        JOIN (
            VALUES {values_sql}
        ) AS v(survey_id, quota_id, country_iso, version)
          ON (ss.survey_id, ss.quota_id, ss.country_iso, ss.version)
           = (v.survey_id, v.quota_id, v.country_iso, v.version);
        """
        params = [item for row in keys for item in row]
        with self.pg_config.make_connection() as conn:
            # self.register_surveystat_key(conn)
            with conn.cursor() as c:
                c.execute(query, params=params)
                res = c.fetchall()
                # print('\n'.join([x['QUERY PLAN'] for x in res]))
        return [SurveyStat.model_validate(x) for x in res]

    def update_surveystats_for_source(
        self,
        source: Source,
        surveys: List[Survey],
        survey_stats: List[SurveyStat],
    ):
        """
        What ym-survey-stats actually calls.
        1. All surveys for this source not in this list of surveys
           get turned off
        2. Get or create all surveys and buyers
        3. Update survey stats
        """
        # Assert the surveys and surveystats we passed are all
        #   for this Source
        survey_source = {s.source for s in surveys}
        assert len(survey_source) == 1 and survey_source == {source}
        # And that the surveys in the surveystats match the passed in Surveys
        surveys_nks = {s.natural_key for s in surveys}
        ss_surveys_nks = {ss.survey_natural_key for ss in survey_stats}
        assert surveys_nks == ss_surveys_nks

        # Turn off not live surveys
        live_surveys = self.survey_manager.filter_by_source_live(source=source)
        live_ids = {s.survey_id for s in live_surveys}
        new_ids = {s.survey_id for s in surveys}
        turn_off_surveys = live_ids - new_ids
        self.survey_manager.turn_off_by_natural_key(
            source=source, survey_ids=turn_off_surveys
        )

        # Create or Update (is_live) Surveys
        res = self.survey_manager.create_or_update(surveys)

        # Update ss
        self.update_or_create(survey_stats=survey_stats)

        return res

    def filter_by_updated_since(self, since):
        return self.filter(updated_after=since, is_live=None)

    def filter_by_live(self):
        return self.filter(is_live=True)

    def make_filter_str(
        self,
        is_live: Optional[bool] = True,
        updated_after: Optional[datetime] = None,
        min_score: Optional[float] = None,
        survey_keys: Optional[Collection[SurveyKey]] = None,
        sources: Optional[Collection[Source]] = None,
        country_iso: Optional[str] = None,
    ):
        filters = []
        params = dict()
        if updated_after is not None:
            params["updated_after"] = updated_after
            filters.append("ss.updated_at >= %(updated_after)s")
        if min_score:
            params["min_score"] = min_score
            filters.append("score >= %(min_score)s")
        if is_live is not None:
            if is_live:
                filters.append("ss.survey_is_live")
            else:
                filters.append("NOT ss.survey_is_live")
        if sources:
            assert survey_keys is None
            params["sources"] = [s.value for s in sources]
            filters.append("survey_source = ANY(%(sources)s)")
        if country_iso:
            params["country_iso"] = country_iso
            filters.append("country_iso = %(country_iso)s")
        if survey_keys is not None:
            # Instead of doing a big IN with a big set of tuples, since we know
            #   we only have N possible sources, we just split by that and do
            #   a set of:
            #      ( (survey_source = 'x' and survey_survey_id IN ('1', '2') ) OR
            #        (survey_source = 'y' and survey_survey_id IN ('3', '4') ) ... )
            sk_filters = []
            survey_source_ids = defaultdict(set)
            for sk in survey_keys:
                source, survey_id = sk.split(":")
                survey_source_ids[Source(source).value].add(survey_id)
            for source, survey_ids in survey_source_ids.items():
                sk_filters.append(
                    f"(survey_source = '{source}' AND survey_survey_id = ANY(%(survey_ids_{source})s))"
                )
                params[f"survey_ids_{source}"] = list(survey_ids)
            # potential bug here ! --v Make sure this is wrapped in parentheses!
            filters.append(f"({' OR '.join(sk_filters)})")

        filter_str = "WHERE " + " AND ".join(filters) if filters else ""
        return filter_str, params

    def filter_count(
        self,
        is_live: Optional[bool] = True,
        updated_after: Optional[datetime] = None,
        min_score: Optional[float] = None,
        survey_keys: Optional[Collection[SurveyKey]] = None,
        sources: Optional[Collection[Source]] = None,
        country_iso: Optional[str] = None,
    ) -> int:
        filter_str, params = self.make_filter_str(
            is_live=is_live,
            updated_after=updated_after,
            min_score=min_score,
            survey_keys=survey_keys,
            sources=sources,
            country_iso=country_iso,
        )
        query = f"""
        SELECT COUNT(1) as cnt
        FROM marketplace_surveystat ss
        {filter_str};
        """
        return self.pg_config.execute_sql_query(query, params=params)[0]["cnt"]

    def filter(
        self,
        is_live: Optional[bool] = True,
        updated_after: Optional[datetime] = None,
        min_score: Optional[float] = None,
        survey_keys: Optional[Collection[SurveyKey]] = None,
        sources: Optional[Collection[Source]] = None,
        country_iso: Optional[str] = None,
        page: Optional[int] = None,
        size: Optional[int] = None,
        order_by: Optional[str] = None,
        debug: Optional[bool] = False,
    ):
        filter_str, params = self.make_filter_str(
            is_live=is_live,
            updated_after=updated_after,
            min_score=min_score,
            survey_keys=survey_keys,
            sources=sources,
            country_iso=country_iso,
        )

        paginated_filter_str = ""
        if page is not None:
            assert page != 0, "page starts at 1"
            size = size if size is not None else 100
            params["offset"] = (page - 1) * size
            params["limit"] = size
            paginated_filter_str = " LIMIT %(limit)s OFFSET %(offset)s"

        order_by_str = ""
        if order_by:
            assert order_by in {"score DESC", "score", "updated_at DESC", "updated_at"}
            order_by_str = f"ORDER BY {order_by}"

        query = f"""
        SELECT
            quota_id, country_iso, cpi,
            complete_too_fast_cutoff,
            prescreen_conv_alpha, prescreen_conv_beta,
            conv_alpha, conv_beta,
            dropoff_alpha, dropoff_beta,
            completion_time_mu, completion_time_sigma,
            mobile_eligible_alpha, mobile_eligible_beta,
            desktop_eligible_alpha, desktop_eligible_beta,
            tablet_eligible_alpha, tablet_eligible_beta,
            long_fail_rate, user_report_coeff, recon_likelihood,
            score_x0, score_x1, updated_at, version, score,
            survey_is_live, survey_source, survey_survey_id
        FROM marketplace_surveystat ss
        {filter_str}
        {order_by_str}
        {paginated_filter_str} ;
        """
        if debug:
            print(query)
            print(params)
        with self.pg_config.make_connection() as conn:
            with conn.cursor() as c:
                c.execute("SET work_mem = '256MB';")
                c.execute("SET statement_timeout = '10s';")
                c.execute(query, params=params)
                res = c.fetchall()
        return [SurveyStat.model_validate(x) for x in res]

    def filter_to_merge_table(
        self,
        is_live: Optional[bool] = True,
        updated_after: Optional[datetime] = None,
        min_score: Optional[float] = 0.0001,
    ):
        survey_stats = self.filter(
            is_live=is_live, updated_after=updated_after, min_score=min_score
        )
        if not survey_stats:
            return None
        extra_cols = {
            "survey_id",
            "quota_id",
            "country_iso",
            "version",
            "updated_at",
            "score_x0",
            "score_x1",
            "survey_is_live",
            "survey_source",
            "survey_survey_id",
        }
        data = []
        for ss in survey_stats:
            d = {k: getattr(ss, v) for k, v in self.SURVEY_STATS_COL_MAP.items()}
            d.update({k: getattr(ss, k) for k in extra_cols})
            d["sid"] = ss.survey_natural_key
            data.append(d)
        df = pd.DataFrame(data)
        df = df.set_index("sid")
        df["cpi"] = df["cpi"].astype(float)
        return df