Criei um ORM em Python para minha empresa · WilianZilv

Aqui na RPS Capital praticamente 100% do nosso dia a dia é lidando com dados no pandas, e muitos desses dados são implementados no SQL de uma maneira que cada projeto seja independente de models definidos via código (sei que é uma má prática, mas precisamos de muita agilidade e pouca complexidade).

Cada consulta que fazemos aqui dentro é única, mesmo puxando dados das mesmas tabelas, o que torna difícil construir funções padronizadas para consultas.

O Problema (ou meu toc)

Muitos aqui ainda fazem consultas da seguinte maneira no Python

import rps_databases # módulo que usamos internamente

db = rps_databases.connect()

SQL = """
    SELECT column_a, column_b
    FROM myschema.mytable
    WHERE id in %s
"""

ids = [1,2,3]

df = db.fetch(SQL, (ids,))

É um exemplo simples, porém, como a query é um texto, eu poderia escrever ela em infinitas variações, que é o que acontece. Eu escrevi do jeito que eu acredito ser o mais legível, mas mesmo assim, não é um padrão.

A solução

df = db.myschema.mytable.get(columns=["column_a", "column_b"], id=ids)

Foi possível resumir em 1 linha o que foi feito em 6 linhas.

O que aconteceu aqui?

A classe Database tem acesso dinâmico às propriedades dela, se eu tento acessar uma propriedade que não existe, ela automaticamente é considerada como um "schema". Isso foi possível utilizando a função "__getattribute__":

class Database:
    ...
    def __getattribute__(self, attr):
        try:
            return object.__getattribute__(self, attr)
        except AttributeError:
            return Schema(db=self, name=attr)

A mesma coisa acontece com a classe Schema, que retorna uma Table ao acessar qualquer propriedade:

class Schema:
    def __init__(self, db, name: str):
        self.db = db
        self.name = name

    def __getattribute__(self, attr):
        try:
            return object.__getattribute__(self, attr)
        except AttributeError:
            return Table(db=self.db, schema=self, name=attr)

Sério, o schema só serve pra isso.

O resto da mágica acontece na classe Table:

class Table:
    def __init__(self, db: "Database", schema: "Schema", name: str):
        self.schema = schema
        self.name = name
        self.db = db

    def path(self):
        return f"{self.schema.name}.{self.name}"

Como ela conhece o caminho do schema e da tabela, só precisamos de uma função para construir a query baseado em alguns requísitos.

    def get(
        self,
        columns: Optional[list] = None,
        distinct: Optional[list] = None,
        min: Optional[list] = None,
        max: Optional[list] = None,
        sum: Optional[list] = None,
        avg: Optional[list] = None,
        **simple_where,
    ) -> pd.DataFrame:
    ...

Com apenas esses parâmetros, já dá pra fazer muita coisa... Aqui vai uns exemplos:

from rps_databases import operators as o

# Idade média por sobrenome onde os usuários são de curitiba
db.public.users.get(columns=["last_name"], avg=["age"], city="Curitiba")

# Idade média por sobrenome onde os usuários não são de curitiba
db.public.users.get(columns=["last_name"], avg=["age"], city=o.Not("Curitiba"))

# Idade média por sobrenome onde os usuários são de curitiba e são paulo
db.public.users.get(columns=["last_name"], avg=["age"], city=["Curitiba", "São Paulo"])

# Idade média por sobrenome onde os usuários não são de curitiba e são paulo
db.public.users.get(columns=["last_name"], avg=["age"], city=o.Not(["Curitiba", "São Paulo"]))

# Data mais recente em uma tabela
db.public.transactions.get(max=["transaction_date"])

# Data mais recente trocando de nome
db.public.transactions.get(max=[
    ("transaction_date", "most_recent_date") # pega o max de transaction_date e renomeia pra most_recent_date
])

# Consulta todas as colunas filtrando as linhas que o last_name contenha em algum lugar "A"
db.public.users.get(last_name=o.Ilike("A"))

Acho que deu pra pegar a ideia, o objetivo foi deixar consultas menos verbosas e menos complexas sem depender de models/schemas definidos.