Rafael Lopes
331a021d9a
Some checks failed
Deploy WiFi-ETL Prod / deploy (push) Failing after 0s
- Adicionado suporte para extração de dados do Ruijie e WiFeed, incluindo autenticação e tratamento de erros. - Adicionado suporte para watermarking em ambas as fontes para extração incremental. - Criado script de transformação para mesclagem de MAC addresses. - Implementado Backfill para WiFeed, permitindo extração histórica com controle de taxa. - Adicionado script de depuração para testes de transformação do WiFeed. - Desenvolvido scripts de implantação e configurações do Docker para setup de produção. - Criado script de inicialização do schema do PostgreSQL em infra/init.sql. - Adicionado teste automatizado para lógica de transformação e carregamento em test_transform_load.py. - Atualizado documentation para implantação e setup de produção.
96 lines
2.5 KiB
Python
96 lines
2.5 KiB
Python
import psycopg2
|
|
import psycopg2.extras
|
|
import logging
|
|
from typing import Optional, Dict, Any
|
|
from contextlib import contextmanager
|
|
|
|
from app.core.config import (
|
|
DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@contextmanager
|
|
def get_connection():
|
|
"""
|
|
Context manager para conexão PostgreSQL.
|
|
Garante fechamento automático e rollback em caso de erro.
|
|
"""
|
|
conn = None
|
|
try:
|
|
conn = psycopg2.connect(
|
|
host=DB_HOST,
|
|
port=DB_PORT,
|
|
dbname=DB_NAME,
|
|
user=DB_USER,
|
|
password=DB_PASSWORD,
|
|
connect_timeout=10
|
|
)
|
|
conn.autocommit = False
|
|
logger.debug("Conexão PostgreSQL estabelecida")
|
|
yield conn
|
|
conn.commit()
|
|
logger.debug("Transação commitada")
|
|
except psycopg2.Error as e:
|
|
if conn:
|
|
conn.rollback()
|
|
logger.error(f"Erro PostgreSQL: {e}")
|
|
raise
|
|
finally:
|
|
if conn:
|
|
conn.close()
|
|
logger.debug("Conexão PostgreSQL fechada")
|
|
|
|
|
|
def execute_query(conn, query: str, params: tuple = None, fetch: bool = False):
|
|
"""
|
|
Executa uma query SQL e opcionalmente retorna resultados.
|
|
|
|
Args:
|
|
conn: Conexão psycopg2
|
|
query: SQL query
|
|
params: Parâmetros para query parametrizada
|
|
fetch: Se True, retorna cursor com resultados
|
|
|
|
Returns:
|
|
cursor (se fetch=True) ou None
|
|
"""
|
|
with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
|
|
cur.execute(query, params)
|
|
if fetch:
|
|
return cur.fetchall()
|
|
return None
|
|
|
|
|
|
def execute_many(conn, query: str, params_list: list):
|
|
"""
|
|
Executa múltiplas inserts/updates em batch.
|
|
|
|
Args:
|
|
conn: Conexão psycopg2
|
|
query: SQL query com placeholders
|
|
params_list: Lista de tuplas de parâmetros
|
|
"""
|
|
with conn.cursor() as cur:
|
|
psycopg2.extras.execute_batch(cur, query, params_list, page_size=1000)
|
|
|
|
|
|
def test_connection() -> bool:
|
|
"""
|
|
Testa conexão com o banco.
|
|
|
|
Returns:
|
|
True se conectado com sucesso
|
|
"""
|
|
try:
|
|
with get_connection() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT 1")
|
|
result = cur.fetchone()
|
|
logger.info(f"Conexão PostgreSQL OK — versão: {conn.server_version}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Falha no teste de conexão: {e}")
|
|
return False
|