Inserir arquivos CSV no banco de dados

Uma das formas de inserir o conteúdo de um arquivo CSV (ou vários) em uma tabela do banco de dados é usando o “adaptador” Psycopg do PostgreSQL no Python. Outra forma é usando o comando “psql” diretamente no terminal Linux. Além disso, é possível inserir linha por linha ou tudo de uma só vez, de apenas um ou de vários arquivos. Serão usados dois métodos: o primeiro usando python e lendo cada um dos arquivos linha por linha (com apensa uma conexão); o segundo, usando psql e inserindo o arquivo CSV diretamente (com apensa uma inserção).

Veja a estrutura da tabela criada para os exemplos:

-- DROP TABLE schema_name.table_name;

CREATE TABLE schema_name.table_name
(
  place character varying(40) NOT NULL,
  date timestamp with time zone NOT NULL,
  var real,
  CONSTRAINT table_name_pkey PRIMARY KEY (place, date)
);

Nesse primeiro exemplo, as queries são escritas em um arquivo .sql, depois executadas em apenas uma conexão com o banco de dados. A primeira query tem o objetivo de limpar o banco de dados, mantendo sua estrutura (TRUNCATE), mas também está comentada a opção para excluir somente as linhas que interessam (DELETE).

Cada query de inserção na tabela (INSERT) é montada a partir de uma linha de um arquivo dentre vários. Ou seja, primeiro o comando “glob” lista todos os nomes de arquivos da pasta “output” que tem o padrão “Out_*.csv”; depois, cada arquivo é lido linha por linha (exceto a primeira, que é do cabeçalho).

#!/usr/bin/python
# -*- coding: utf-8 -*-

import psycopg2
import os
import sys
import glob

# Caminho completo do diretório do script
path = os.path.dirname(os.path.realpath(__file__))
nome_da_tabela = table_name

# Abrir arquivo de saída para gravação das queries
file_sql = "%s/csv2db.sql" %path
fileout = open(file_sql,'w')

# Escrever query para apagar conteúdo da tabela (mas não estrutura)
querydel = "TRUNCATE table %s;\n" % nome_da_tabela
fileout.write(querydel)

# Listar arquivos e obter os nomes das cidades
mydir = "%s/output/Out_*.csv" %path
listing = sorted(glob.glob(mydir))
for filename in listing:
	temp = filename.split(os.sep)[-1]
	nome_cidade = temp.replace("Out_","").replace(".csv","")
	with open('%s/output/Out_%s.csv' %(path, nome_cidade)) as arq:
		temp = arq.readline()
		cabecalho = temp.strip().split(",")
		for linha in arq:
			linha_array = linha.strip().split(",")
			# Montar query
			#querydel = "DELETE FROM %s WHERE data = '%s' AND cidade = '%s';\n" %(nome_da_tabela, linha_array[0], nome_cidade)
			query = "INSERT INTO %s(place,date,var) VALUES ('%s','%s',%s);\n" %(nome_da_tabela, nome_cidade, linha_array[0], linha_array[1])
			# Imprimir query em arquivo SQL
			#fileout.write(querydel)
			fileout.write(query)
fileout.close()

# Abrir conexão, executar SQL e fechar conexão
conn = psycopg2.connect(host = "HOST_NAME", database = "DB_NAME", user="USER_NAME")
curs = conn.cursor()
curs.execute(open(file_sql, "r").read())
# Fechar ponteiro e conexão
curs.close()
conn.close()
# Fechar conexão se fizer inserção em banco
conn.commit()

Uma desvantagem desse método é a quantidade grande de inserções realizadas na base de dados.

Já no segundo exemplo, é feita uma inserção de um arquivo CSV mais rápida e limpa usando a estrutura “COPY … FROM” do SQL. Como não há necessidade de ler o arquivo linha por linha, é possível fazer tudo com apenas uma linha de comando diretamente no terminal, usando o programa “psql”.

Como são vários arquivos contendo um cabeçalho e as linhas com os dados, primeiro será criado um arquivo temporário com o cabeçalho e todas as outras linhas (exceto a primeira de cada arquivo):

echo "place,date,var" > output/Out_todos.csv
awk FNR-1 output/Out_*.csv >> output/Out_todos.csv

O comando do SQL para inserir diretamente um arquivo CSV em uma tabela tem a seguinte sintaxe:

COPY schema_name.table_name FROM '/complete_path/Filename.csv';

No entanto, se não for executado pelo super usuário, pode aparecer o seguinte erro:

psycopg2.ProgrammingError: must be superuser to COPY to or from a file
HINT:  Anyone can COPY to stdout or from stdin. psql's \copy command also works for anyone.

Seguindo a dica, é possível jogar a saída de um programa de visualização de arquivos (o cat, por exemplo) para a entrada (stdin) do comando psql usando um redirecionador de fluxo (o pipe). O comando recebe como parâmetros o endereço da base de dados (-H), o usuário (-U), o nome da tabela (-W), uma indicação de que não tem senha para acessar o banco e o comando (-c):

cat output/Out_todos.csv | psql -h HOST_NAME -U USER_NAME -W table_name --no-password -c "TRUNCATE table table_name; COPY schema_name.tabel_name FROM STDIN WITH CSV HEADER NULL AS '' DELIMITER ',';"

Note que o comando mudou um pouco se comparado à versão original. Primeiro, foi incluída uma query para apagar o conteúdo da tabela. Segundo, foi alterada a entrada da query para receber a stdin em vez do arquivo, que também indica que contém um cabeçalho, que os campos nulos estão vazios (isso evita erros do tipo “invalid input syntax for type timestamp”) e que tudo está delimitado por vírgula.

Como foram definidas as colunas “place” e “date” como chaves primárias, caso tenha algum registro duplicado, não será realizada NENHUMA mudança.