<commit_before>#!/usr/bin/python
import psycopg2 as pg2
import pandas.io.sql as psql
import pandas as pd
import os
import shutil
import glob
import datetime
from collections import defaultdict
from dateutil.parser import parse

from lode.utilities.util import (parse_date, load_config,
                                 get_file_year_str, meta_path)

from lode.database.helpers import (check_csv_headers,
                                   strip_fileendings)

from itertools import izip
from OfferPandas import Frame


class NZEMDB(object):
    """docstring for NZEMDB"""
    def __init__(self, DB_KEY):
        super(NZEMDB, self).__init__()

        self.refresh_config()

        self.db_key = DB_KEY
        self.schemas = self.CONFIG[DB_KEY]["schemas"]

        self.db = self.CONFIG[DB_KEY]["database_name"]
        self.user = self.CONFIG[DB_KEY]["database_user"]
        self.host = "localhost"
        self.password = self.CONFIG[DB_KEY]["database_pass"]
        self.conn_string = """dbname=%s user=%s password=%s
                              host=%s""" % (self.db, self.user, self.password,
                                            self.host)

    def refresh_config(self):
        """ This permits hot loading of the config file instead of linking
        it to only be initialised on startup
        """
        self.CONFIG = load_config()

        return self

    def insert_to_database(self, query, csvfile, tabname):
        try:
            self.execute_and_commit_sql(query)
        except pg2.IntegrityError as e:
            print "You're trying to add rows which already exist!"
            print e

        except pg2.ProgrammingError as e:
            print e
            #print "Going to do some permission editing from a local drive"
            #print "You're probably running from an external drive if you see"
            #print "This often"
            fName = os.path.basename(csvfile)
            homeName = os.path.expanduser("~/%s" % fName)
            shutil.copy(csvfile, homeName)
            os.chmod(homeName, 0644)
            query = """COPY %s FROM '%s' DELIMITER ',' CSV HEADER;
                    """ % (tabname, homeName)
            try:
                self.execute_and_commit_sql(query)
            except pg2.IntegrityError as e:
                print "You're trying to add rows which already exist!"
                print e
            except pg2.DataError as e:
                print e
                strip_fileendings(homeName)
                self.execute_and_commit_sql(query)

            os.remove(homeName)

    def insert_from_csv(self, table, csvfile):

        # Check the Schemas for the split year:
        if self.schemas[table]['split_by_year']:
            year = get_file_year_str(csvfile)
            table_name = "%s_%s" % (table, year)
        else:
            table_name = table

        table_headers = self.get_column_names(table_name)
        insert_headers = ",".join([x[0] for x in
                                   table_headers if "key" not in x[0]])
        tabname = "%s(%s)" % (table_name, insert_headers)

        existing_headers = check_csv_headers(csvfile, table_headers)

        if existing_headers:
            query = """COPY %s FROM '%s' DELIMITER ',' CSV HEADER;
                    """ % (tabname, csvfile)
        else:
            query = """COPY %s FROM '%s' DELIMITER ',' CSV;
                    """ % (tabname, csvfile)

        self.insert_to_database(query, csvfile, tabname)

    def strip_fileendings(self, fName):
        print "Attempting to strip the shitty endings"
        with open(fName, 'rb') as f:
            data = f.readlines()

        data_new = [d.replace("\r\n", "\n") for d in data]

        with open(fName, 'wb') as f:
            for row in data_new:
                f.write(row)

    def execute_and_commit_sql(self, sql):

        with pg2.connect(self.conn_string) as conn:
            with conn.cursor() as curs:
                curs.execute(sql)
                conn.commit()

    def ex_sql_and_fetch(self, sql):
        with pg2.connect(self.conn_string) as conn:
            with conn.cursor() as curs:
                curs.execute(sql)
                records = curs.fetchall()

        return records

    def get_column_names(self, table):

        sql = """SELECT column_name FROM information_schema.columns
                 WHERE table_schema='public' AND table_name='%s'""" % table

        return self.ex_sql_and_fetch(sql)[::-1]

    def create_all_tables(self):

        for key in self.schemas.keys():
            with open(self.schemas[key]["schema_location"], 'rb') as f:
                sql = f.read()

            sql = """%s""" % sql
            if self.schemas[key]["split_by_year"]:
                for year in self.schemas[key]["split_years"]:
                    sql_year = sql % year
                    self.execute_and_commit_sql(sql_year)
            else:
                self.execute_and_commit_sql(sql)

    def drop_table(self, table):
        sql = """DROP TABLE %s""" % table
        self.execute_and_commit_sql(sql)

    def drop_tables(self, master_table):

        if self.schemas[master_table]["split_by_year"]:
            for year in self.schemas[master_table]["split_years"]:
                self.drop_table("_".join([master_table, str(year)]))

    def query_to_df(self, sql):
        """ Use a Generator here as we're Lazy """

        with pg2.connect(self.conn_string) as conn:
            yield psql.read_sql(sql, conn)

    def insert_many_csv(self, table, folder):

        allcsv_files = glob.glob(folder + "/*.csv")
        for f in allcsv_files:
            print "Attempting to load %s" % f
            self.insert_from_csv(table, f)
            print "%s succesfully loaded to %s" % (f, table)

    def list_all_tables(self):
        return self.ex_sql_and_fetch("SELECT * FROM pg_catalog.pg_tables")<commit_msg>Fixing a generator which was incorrectly returning connections<commit_after>#!/usr/bin/python
import psycopg2 as pg2
import pandas.io.sql as psql
import pandas as pd
import os
import shutil
import glob
import datetime
from collections import defaultdict
from dateutil.parser import parse

from lode.utilities.util import (parse_date, load_config,
                                 get_file_year_str, meta_path)

from lode.database.helpers import (check_csv_headers,
                                   strip_fileendings)

from itertools import izip
from OfferPandas import Frame


class NZEMDB(object):
    """docstring for NZEMDB"""
    def __init__(self, DB_KEY):
        super(NZEMDB, self).__init__()

        self.refresh_config()

        self.db_key = DB_KEY
        self.schemas = self.CONFIG[DB_KEY]["schemas"]

        self.db = self.CONFIG[DB_KEY]["database_name"]
        self.user = self.CONFIG[DB_KEY]["database_user"]
        self.host = "localhost"
        self.password = self.CONFIG[DB_KEY]["database_pass"]
        self.conn_string = """dbname=%s user=%s password=%s
                              host=%s""" % (self.db, self.user, self.password,
                                            self.host)

    def refresh_config(self):
        """ This permits hot loading of the config file instead of linking
        it to only be initialised on startup
        """
        self.CONFIG = load_config()

        return self

    def insert_to_database(self, query, csvfile, tabname):
        try:
            self.execute_and_commit_sql(query)
        except pg2.IntegrityError as e:
            print "You're trying to add rows which already exist!"
            print e

        except pg2.ProgrammingError as e:
            print e
            #print "Going to do some permission editing from a local drive"
            #print "You're probably running from an external drive if you see"
            #print "This often"
            fName = os.path.basename(csvfile)
            homeName = os.path.expanduser("~/%s" % fName)
            shutil.copy(csvfile, homeName)
            os.chmod(homeName, 0644)
            query = """COPY %s FROM '%s' DELIMITER ',' CSV HEADER;
                    """ % (tabname, homeName)
            try:
                self.execute_and_commit_sql(query)
            except pg2.IntegrityError as e:
                print "You're trying to add rows which already exist!"
                print e
            except pg2.DataError as e:
                print e
                strip_fileendings(homeName)
                self.execute_and_commit_sql(query)

            os.remove(homeName)

    def insert_from_csv(self, table, csvfile):

        # Check the Schemas for the split year:
        if self.schemas[table]['split_by_year']:
            year = get_file_year_str(csvfile)
            table_name = "%s_%s" % (table, year)
        else:
            table_name = table

        table_headers = self.get_column_names(table_name)
        insert_headers = ",".join([x[0] for x in
                                   table_headers if "key" not in x[0]])
        tabname = "%s(%s)" % (table_name, insert_headers)

        existing_headers = check_csv_headers(csvfile, table_headers)

        if existing_headers:
            query = """COPY %s FROM '%s' DELIMITER ',' CSV HEADER;
                    """ % (tabname, csvfile)
        else:
            query = """COPY %s FROM '%s' DELIMITER ',' CSV;
                    """ % (tabname, csvfile)

        self.insert_to_database(query, csvfile, tabname)

    def strip_fileendings(self, fName):
        print "Attempting to strip the shitty endings"
        with open(fName, 'rb') as f:
            data = f.readlines()

        data_new = [d.replace("\r\n", "\n") for d in data]

        with open(fName, 'wb') as f:
            for row in data_new:
                f.write(row)

    def execute_and_commit_sql(self, sql):

        with pg2.connect(self.conn_string) as conn:
            with conn.cursor() as curs:
                curs.execute(sql)
                conn.commit()

    def ex_sql_and_fetch(self, sql):
        with pg2.connect(self.conn_string) as conn:
            with conn.cursor() as curs:
                curs.execute(sql)
                records = curs.fetchall()

        return records

    def get_column_names(self, table):

        sql = """SELECT column_name FROM information_schema.columns
                 WHERE table_schema='public' AND table_name='%s'""" % table

        return self.ex_sql_and_fetch(sql)[::-1]

    def create_all_tables(self):

        for key in self.schemas.keys():
            with open(self.schemas[key]["schema_location"], 'rb') as f:
                sql = f.read()

            sql = """%s""" % sql
            if self.schemas[key]["split_by_year"]:
                for year in self.schemas[key]["split_years"]:
                    sql_year = sql % year
                    self.execute_and_commit_sql(sql_year)
            else:
                self.execute_and_commit_sql(sql)

    def drop_table(self, table):
        sql = """DROP TABLE %s""" % table
        self.execute_and_commit_sql(sql)

    def drop_tables(self, master_table):

        if self.schemas[master_table]["split_by_year"]:
            for year in self.schemas[master_table]["split_years"]:
                self.drop_table("_".join([master_table, str(year)]))

    def query_to_df(self, sql):
        """ Use a Generator here as we're Lazy """

        with pg2.connect(self.conn_string) as conn:
            return psql.read_sql(sql, conn)

    def insert_many_csv(self, table, folder):

        allcsv_files = glob.glob(folder + "/*.csv")
        for f in allcsv_files:
            print "Attempting to load %s" % f
            self.insert_from_csv(table, f)
            print "%s succesfully loaded to %s" % (f, table)

    def list_all_tables(self):
        return self.ex_sql_and_fetch("SELECT * FROM pg_catalog.pg_tables")
{'fraction_non_alphanumeric': 0.10167201486235433, 'fraction_numerical': 0.0027022462421888195, 'mean_word_length': 1.9696589769307924, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4857145', 'n_tokens_mistral': 3602, 'n_tokens_neox': 3396, 'n_words': 898}