Function load_table_counts

tests/pipeline/utils.py:412–441

Return number of rows for each requested table name. Args: p (dlt.Pipeline): Pipeline instance. *table_names (str): Optional list of table names. When omitted, counts for all data tables in the default schema are returned. Returns: DictStrAny: Mapping of

(p: dlt.Pipeline, *table_names: str)

def load_table_counts(p: dlt.Pipeline, *table_names: str) -> DictStrAny:
    """Return number of rows for each requested table name.

    Args:
        p (dlt.Pipeline): Pipeline instance.
        *table_names (str): Optional list of table names. When omitted, counts
            for all data tables in the default schema are returned.

    Returns:
        DictStrAny: Mapping of table name to integer row count.
    """
    if not table_names:
        table_names = [table["name"] for table in p.default_schema.data_tables()]  # type: ignore[assignment]

    # filesystem with sftp requires a fallback
    if _is_sftp(p):
        file_tables = _load_tables_to_dicts_fs(p, *table_names)
        return {table_name: len(items) for table_name, items in file_tables.items()}

    # NOTE: filesystem with abfss and no table format requires a fallback where we get each table count individually
    # this seems to be a bug in duckdb abfss and might be resolved in a future version.
    if _is_abfss(p):
        table_counts = {}
        for table in table_names:
            table_counts[table] = p.dataset().row_counts(table_names=[table]).fetchall()[0][1]
        return table_counts

    # otherwise we can use the dataset row counts
    counts = p.dataset().row_counts(table_names=list(table_names)).fetchall()
    return {row[0]: row[1] for row in counts}


def assert_empty_tables(p: dlt.Pipeline, *table_names: str) -> None: