Cuando preparamos una clase o una demo, lo último que queremos es perder la atención de nuestra audiencia a los 3 minutos. Tenemos que encontrar la manera para tener entornos de pruebas - ya sea para enseñar o testear - que sean fáciles de usar y reproducibles en el tiempo.

Por suerte, hay bases de datos que nos permiten distribuir imágenes de Docker que ya contengan los ingredientes que necesitas para tu próxima charla (o benchmark)!

Veamos como podemos conseguirlo sin dolores de cabeza con Postgres.

1. Preparar los datos

Para empezar necesitamos un fichero SQL que se encargue de:

1. Crear el usuario y contraseña que usaremos para conectarnos a la DB, ya que no necesariamente queremos usar los usuarios administrativos por defecto.

2. Definir el modelo de datos (esquemas, tablas,…)

Un ejemplo sencillo en Postgres podría ser:

CREATE USER demo WITH PASSWORD ‘password’;
GRANT ALL PRIVILEGES ON DATABASE postgres TO demo;

-- CREATE TABLE...

En este caso, vamos a escribir un script algo más interesante que nos ayudará a generar datos para una tabla EMPLOYEES. Sí, hay librerías como Faker en Python que nos ayudan una barbaridad y tienen funciones super sencillas para generar nombres, direcciones, textos,… Aún así, tener herramientas a mano para poder jugar directamente en una DB puede resultar útil.

No siempre será todo SQL o todo Python. La herramienta o el lenguaje es lo de menos, aprendamos a usar bien todo lo que tenemos a mano.

CREATE TABLE EMPLOYEES (
  id SERIAL PRIMARY KEY,
  first_name VARCHAR(50),
  last_name VARCHAR(50),
  email VARCHAR(50),
  mobile_no BIGINT,
  date_of_birth DATE
);

ALTER TABLE public.employees OWNER TO demo;

CREATE FUNCTION get_random_string() RETURNS TEXT LANGUAGE SQL AS $$
SELECT STRING_AGG (SUBSTR(’abcdefghijklmnopqrstuvwxyz’, CEIL(RANDOM() * 26)::integer, 1), ‘’)
FROM GENERATE_SERIES(1, 10)
$$;

CREATE FUNCTION insert_record() RETURNS VOID LANGUAGE PLPGSQL AS $$
DECLARE first_name TEXT= INITCAP(get_random_string());
DECLARE last_name TEXT= INITCAP(get_random_string());
DECLARE email TEXT= LOWER(CONCAT(first_name, ‘.’, last_name, ‘@gmail.com’));
DECLARE mobile_no BIGINT=CAST(1000000000 + FLOOR(RANDOM() * 9000000000) AS BIGINT);
DECLARE date_of_birth DATE= CAST(NOW() - INTERVAL ‘100 year’ * RANDOM() AS DATE);
BEGIN
INSERT INTO EMPLOYEES (
  first_name,
  last_name,
  email,
  mobile_no,
  date_of_birth
) VALUES (first_name, last_name, email, mobile_no, date_of_birth);
END;
$$;

SELECT insert_record() FROM GENERATE_SERIES(1, 1000000);

Este script nos ayuda a generar INSERTs mediante funciones definidas directamente en Postgres, con strings, números y fechas aleatorias. Gracias al GENERATE_SERIES podemos fácilmente ejecutar la función insert_record 1M de veces, con lo que la tabla EMPLOYEES acabará con 1M de registros.

Además, si no conocíais la tabla pg_stat_statements, os recomiendo habilitarla para poder revisar estadísticas de ejecución de queries y planes de ejecución:

CREATE EXTENSION pg_stat_statements;
GRANT pg_read_all_stats TO demo;

Cerramos el primer capítulo guardando nuestro fichero SQL con el nombre como queramos, por ejemplo postgres-script.sql.

2. Crear el Dockerfile

Ahora toca asegurarnos que el script se va a cargar cuando se genere la imagen de nuestra DB:

FROM postgres:16
WORKDIR /docker-entrypoint-initdb.d
COPY ./postgres-script.sql .
RUN chmod -R 775 /docker-entrypoint-initdb.d

Postgres tiene un directorio /docker-entrypoint-initdb.d donde si ponemos ficheros de SQL, van a ejecutarse cuando la DB arranque. Copiando nuestro script ahí nos aseguramos que cuando alguien use nuestra imagen, tenga el usuario y los datos que hemos definido.

3. Construir y ejecutar

Vamos a crear la imagen con:

docker build . -t demo/postgres:local

y la ejecutamos mapeando bien los puertos:

docker run \
  -p 5432:5432 \
  -e POSTGRES_PASSWORD=password \
  -e POSTGRES_USER=postgres \
  --name demo_pg \
  --rm \
  demo/postgres:local postgres \
  -c shared_preload_libraries=pg_stat_statements \
  -c pg_stat_statements.track=all \
  -c max_connections=200

Ya sólo queda conectarnos a localhost:5432 (yo uso DBeaver) y estamos listos!

Un abrazo.

Esta librería nos proporciona abstracciones que son realmente cómodas para que no tengamos que pensar demasiado qué pasa en el lado de la DB. Sacamos ciertos datos fuera, los mapeamos a objetos en Python y seguimos con nuestras vidas. Aun así, debemos ser capaces de elegir correctamente como interactuamos con el origen de los datos.

Uno de los puntos principales a tener en cuenta es administrar correctamente la memoria de nuestros procesos: Si trabajamos con muchos datos y los intentamos leer todos de golpe, seguramente acabemos con un Out of Memory Error, que se puede traducir en un error 137 (SIGKILL) si estamos trabajando en entornos containerizados.

En este post vamos a analizar tres estrategias diferentes, ver cuánta memoria consumen y como se comportan en el lado de la DB, leyendo siempre los mismos datos de la misma tabla origen.

Nota: Para el análisis vamos a usar una única tabla con 1M de registros en Postgres. Para el perfilado de la memoria, usaremos mprof.

all()

from sqlalchemy import create_engine, text

connection_string = "postgresql://demo:password@localhost:5432/postgres"
engine = create_engine(connection_string)

with engine.connect() as conn:
    res = conn.execute(text("SELECT * FROM public.employees")).all()
    counter = 0
    for elem in res:
        counter += 1

print(counter)

Vamos a empezar con la forma más simple: Una única query a la tabla employees que va a intentar sacar el millón de registros de una sentada. El proceso corre deprisa, pero si miramos cuánta memoria nos estamos comiendo, llega por encima de los 600MB.

Esto no es necesariamente malo, sólo tenemos que entender los recursos que debemos proveer al servicio, o como se va a comportar si los datos evolucionan en el futuro.

Entender como el cliente de Python gestiona este escenario está bien, pero es aún más interesante si entramos directamente en como la DB está respondiendo a nuestra llamada, con la siguiente query:

SELECT
  u.usename,
  stats.queryid,
  stats.query,
  stats.total_exec_time
FROM pg_stat_statements stats
join pg_catalog.pg_user u
  on u.usesysid = stats.userid
where u.usename = 'demo';

Aquí no debería ser una sorpresa ver una única entrada con la query SELECT * FROM public.employees. Estamos levantando un único Result de la DB equivalente a toda la tabla, e iteramos sobre este resultado directamente en el cliente.

Veamos si lo podemos hacer algo mejor.

yield_per()

with engine.connect() as conn:
    res = conn.execute(text(
      "SELECT * FROM public.employees"
    )).yield_per(10)
    counter = 0
    for rows in res:
        counter += 1
print(counter)

Fijaos como ahora estamos llamando yield_per sobre el objeto que nos proporciona el execute(...).

Esto crea un cursor en el cliente con el siguiente perfil de memoria:

Ahora estamos justo por debajo de 250MB, bastante mejor que con la llamada a all(). Si revisamos de nuevo las entradas en pg_stats_statements, todavía veremos una única llamada a la query SELECT * FROM public.employees.

En este escenario seguimos haciendo una única llamada a la DB, pero estamos iterando sobre los datos de forma algo más eficiente del lado del cliente.

Todavía podemos hacerlo mejor.

Cursores en el servidor

Según la docu de Postgres: En lugar de ejecutar una consulta completa de una vez, es posible configurar un cursor que encapsule la consulta y luego leer el resultado unas pocas filas a la vez.

Esta sería una estrategia estupenda a seguir aquí. En vez de cargar toda la tabla employees, podríamos obtener unas pocas filas en cada llamada. Esto implica múltiples viajes a la DB, pero el uso de memoria será mucho menor.

Por suerte, SQLAlchemy nos permite hacer streams de resultados para usar directamente los cursores de las bases de datos. Iremos algo más lentos, pero la memoria del cliente no va a explotar.

with engine.connect() as conn:
    result = conn.execution_options(stream_results=True, max_row_buffer=100).execute(
        text("select * from public.employees")
    )
    counter = 0
    for row in result:
        counter += 1
print(counter)

Vamos a preparar el execution_options antes de llamar a execute, tuneando la comunicación db <> cliente.

Esto nos deja un uso de memoria algo mayor a 40MB!

Si ahora revisamos las entradas en pg_stats_statements, veremos que realmente estamos levantando la tabla employees a trozos (en chunks):

DECLARE "c_10289e140_1" CURSOR WITHOUT HOLD FOR select * from public.employees
FETCH FORWARD 100 FROM "c_10289e140_1"
FETCH FORWARD 25 FROM "c_10289e140_1"
...

Conclusión

Qué método deberíamos usar? Como siempre, depende: del tamaño de los datos, la velocidad que necesita el proceso, los recursos con los que contamos, etc. Lo importante es elegir con conocimiento de causa, una vez entendemos las diferencias de como nuestro proceso se va a comportar.

En el siguiente post, vamos a adentrarnos en cómo preparar el set de datos con el que hemos jugado de una manera supercómoda.

Un abrazo.