<commit_before>import json
import os
from os import path

from datetime import datetime
from psycopg2 import connect
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT


class Etl(object):
    def __init__(self, connection_string, folder, wipe=False):
        self.wipe = wipe
        self.connection_string = connection_string
        self.folder = folder

    def create_tables(self):
        with connect(self.connection_string) as con:
            con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
            with con.cursor() as cur:
                query = """
                    DROP TABLE IF EXISTS connection CASCADE;
                    CREATE TABLE connection(
                      departureStop VARCHAR(255) NOT NULL,
                      departureTime TIMESTAMP NOT NULL,
                      departureDate VARCHAR(255) NOT NULL,
                      departureDelay FLOAT NOT NULL,
                      arrivalStop VARCHAR(255) NOT NULL,
                      arrivalTime TIMESTAMP NOT NULL,
                      arrivalDate VARCHAR(255) NOT NULL,
                      arrivalDelay FLOAT NOT NULL,
                      route VARCHAR(255) NOT NULL,
                      trip VARCHAR(255) NOT NULL,
                      id VARCHAR(255) NOT NULL PRIMARY KEY
                    );

                    DROP TABLE IF EXISTS station CASCADE;
                    CREATE TABLE station(
                      name VARCHAR(255) NOT NULL,
                      country VARCHAR(255) NOT NULL,
                      longitude FLOAT NOT NULL,
                      latitude FLOAT NOT NULL,
                      id VARCHAR(255) NOT NULL PRIMARY KEY
                    );

                    DROP TABLE IF EXISTS occupancy CASCADE;
                    CREATE TABLE occupancy(
                      timestamp TIMESTAMP NOT NULL,
                      connection VARCHAR(255) NOT NULL,
                      stationFrom VARCHAR(255) NOT NULL,
                      stationTo VARCHAR(255) NOT NULL,
                      date VARCHAR(255) NOT NULL,
                      vehicle VARCHAR(255) NOT NULL,
                      occupancy VARCHAR(255) NOT NULL,
                      userAgent VARCHAR(255) NOT NULL
                    );

                    DROP TABLE IF EXISTS distance CASCADE;
                    CREATE TABLE distance(
                      stationFrom VARCHAR(255) NOT NULL,
                      stationTo VARCHAR(255) NOT NULL,
                      trip VARCHAR(255) NOT NULL,
                      date VARCHAR(255) NOT NULL,
                      distance INT NOT NULL
                    );

                    DROP VIEW IF EXISTS events;
                    CREATE VIEW events AS (
                      SELECT
                        c.*,
                        c.departuretime AT TIME ZONE 'CET'                    AS departuretime_local,
                        c.arrivaltime AT TIME ZONE 'CET'                      AS arrivaltime_local,
                        sa.name                                               AS arrival,
                        sa.longitude                                          AS arrival_longitude,
                        sa.latitude                                           AS arrival_latitude,
                        sa.country                                            AS arrival_country,
                        sa.longitude :: TEXT                                  AS arrival_longitude_str,
                        sa.latitude :: TEXT                                   AS arrival_latitude_str,
                        EXTRACT(MONTH FROM c.departuretime)                   AS departure_month,
                        EXTRACT(DAY FROM c.departuretime)                     AS departure_day,
                        EXTRACT(DOW FROM c.departuretime)                     AS departure_dow,
                        EXTRACT(HOUR FROM c.departuretime AT TIME ZONE 'CET') AS departure_hour,
                        sd.name                                               AS departure,
                        sd.longitude                                          AS departure_longitude,
                        sd.latitude                                           AS departure_latitude,
                        sd.country                                            AS departure_country,
                        sd.longitude :: TEXT                                  AS departure_longitude_str,
                        sd.latitude :: TEXT                                   AS departure_latitude_str,
                        regexp_replace(c.route, '[^A-Z]', '', 'g')            AS traintype

                      FROM connection c INNER JOIN station sa ON c.arrivalstop = sa.id
                        INNER JOIN station sd ON c.departurestop = sd.id);
                """
                # query= """
                # DROP TABLE IF EXISTS distance;
                #     CREATE TABLE distance(
                #       stationFrom VARCHAR(255) NOT NULL,
                #       stationTo VARCHAR(255) NOT NULL,
                #       vehicle VARCHAR(255) NOT NULL,
                #       date VARCHAR(255) NOT NULL,
                #       distance INT NOT NULL
                #     );
                # """
                cur.execute(query)
                cur.close()

    def extract_connection_files(self, folder, batch_size, last_departure_time=None):
        json_files = [pos_json for pos_json in os.listdir(folder) if pos_json.endswith('.json')].sort()
        result = []
        last_js = ''
        for js in json_files:
            last_js = js
            ts = datetime.strptime(js.replace(".json", ""), "%Y-%m-%dT%H-%M")
            if last_departure_time is not None and ts < last_departure_time:
                continue
            with open(os.path.join(folder, js)) as json_file:
                data = json.load(json_file)
                for conn in data['@graph']:
                    result.append([
                        conn['departureStop'].replace('http://irail.be/stations/NMBS/', ''),
                        conn['departureTime'],
                        conn['departureTime'][:10].replace('-', ''),
                        conn['departureDelay'] if 'departureDelay' in conn else "0",
                        conn['arrivalStop'].replace('http://irail.be/stations/NMBS/', ''),
                        conn['arrivalTime'],
                        conn['arrivalTime'][:10].replace('-', ''),
                        conn['arrivalDelay'] if 'arrivalDelay' in conn else "0",
                        conn['gtfs:route'].replace('http://irail.be/vehicle/', ''),
                        conn['gtfs:trip'].replace('http://irail.be/trips/', ''),
                        conn['@id'].replace('http://irail.be/connections/', '')])
                if len(result) > batch_size:
                    return_value = result
                    result = []<commit_msg>make sure connection files are in right order<commit_after>import json
import os
from os import path

from datetime import datetime
from psycopg2 import connect
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT


class Etl(object):
    def __init__(self, connection_string, folder, wipe=False):
        self.wipe = wipe
        self.connection_string = connection_string
        self.folder = folder

    def create_tables(self):
        with connect(self.connection_string) as con:
            con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
            with con.cursor() as cur:
                query = """
                    DROP TABLE IF EXISTS connection CASCADE;
                    CREATE TABLE connection(
                      departureStop VARCHAR(255) NOT NULL,
                      departureTime TIMESTAMP NOT NULL,
                      departureDate VARCHAR(255) NOT NULL,
                      departureDelay FLOAT NOT NULL,
                      arrivalStop VARCHAR(255) NOT NULL,
                      arrivalTime TIMESTAMP NOT NULL,
                      arrivalDate VARCHAR(255) NOT NULL,
                      arrivalDelay FLOAT NOT NULL,
                      route VARCHAR(255) NOT NULL,
                      trip VARCHAR(255) NOT NULL,
                      id VARCHAR(255) NOT NULL PRIMARY KEY
                    );

                    DROP TABLE IF EXISTS station CASCADE;
                    CREATE TABLE station(
                      name VARCHAR(255) NOT NULL,
                      country VARCHAR(255) NOT NULL,
                      longitude FLOAT NOT NULL,
                      latitude FLOAT NOT NULL,
                      id VARCHAR(255) NOT NULL PRIMARY KEY
                    );

                    DROP TABLE IF EXISTS occupancy CASCADE;
                    CREATE TABLE occupancy(
                      timestamp TIMESTAMP NOT NULL,
                      connection VARCHAR(255) NOT NULL,
                      stationFrom VARCHAR(255) NOT NULL,
                      stationTo VARCHAR(255) NOT NULL,
                      date VARCHAR(255) NOT NULL,
                      vehicle VARCHAR(255) NOT NULL,
                      occupancy VARCHAR(255) NOT NULL,
                      userAgent VARCHAR(255) NOT NULL
                    );

                    DROP TABLE IF EXISTS distance CASCADE;
                    CREATE TABLE distance(
                      stationFrom VARCHAR(255) NOT NULL,
                      stationTo VARCHAR(255) NOT NULL,
                      trip VARCHAR(255) NOT NULL,
                      date VARCHAR(255) NOT NULL,
                      distance INT NOT NULL
                    );

                    DROP VIEW IF EXISTS events;
                    CREATE VIEW events AS (
                      SELECT
                        c.*,
                        c.departuretime AT TIME ZONE 'CET'                    AS departuretime_local,
                        c.arrivaltime AT TIME ZONE 'CET'                      AS arrivaltime_local,
                        sa.name                                               AS arrival,
                        sa.longitude                                          AS arrival_longitude,
                        sa.latitude                                           AS arrival_latitude,
                        sa.country                                            AS arrival_country,
                        sa.longitude :: TEXT                                  AS arrival_longitude_str,
                        sa.latitude :: TEXT                                   AS arrival_latitude_str,
                        EXTRACT(MONTH FROM c.departuretime)                   AS departure_month,
                        EXTRACT(DAY FROM c.departuretime)                     AS departure_day,
                        EXTRACT(DOW FROM c.departuretime)                     AS departure_dow,
                        EXTRACT(HOUR FROM c.departuretime AT TIME ZONE 'CET') AS departure_hour,
                        sd.name                                               AS departure,
                        sd.longitude                                          AS departure_longitude,
                        sd.latitude                                           AS departure_latitude,
                        sd.country                                            AS departure_country,
                        sd.longitude :: TEXT                                  AS departure_longitude_str,
                        sd.latitude :: TEXT                                   AS departure_latitude_str,
                        regexp_replace(c.route, '[^A-Z]', '', 'g')            AS traintype

                      FROM connection c INNER JOIN station sa ON c.arrivalstop = sa.id
                        INNER JOIN station sd ON c.departurestop = sd.id);
                """
                # query= """
                # DROP TABLE IF EXISTS distance;
                #     CREATE TABLE distance(
                #       stationFrom VARCHAR(255) NOT NULL,
                #       stationTo VARCHAR(255) NOT NULL,
                #       vehicle VARCHAR(255) NOT NULL,
                #       date VARCHAR(255) NOT NULL,
                #       distance INT NOT NULL
                #     );
                # """
                cur.execute(query)
                cur.close()

    def extract_connection_files(self, folder, batch_size, last_departure_time=None):
        json_files = [pos_json for pos_json in os.listdir(folder) if pos_json.endswith('.json')]
        json_files.sort()
        result = []
        last_js = ''
        for js in json_files:
            last_js = js
            ts = datetime.strptime(js.replace(".json", ""), "%Y-%m-%dT%H-%M")
            if last_departure_time is not None and ts < last_departure_time:
                continue
            with open(os.path.join(folder, js)) as json_file:
                data = json.load(json_file)
                for conn in data['@graph']:
                    result.append([
                        conn['departureStop'].replace('http://irail.be/stations/NMBS/', ''),
                        conn['departureTime'],
                        conn['departureTime'][:10].replace('-', ''),
                        conn['departureDelay'] if 'departureDelay' in conn else "0",
                        conn['arrivalStop'].replace('http://irail.be/stations/NMBS/', ''),
                        conn['arrivalTime'],
                        conn['arrivalTime'][:10].replace('-', ''),
                        conn['arrivalDelay'] if 'arrivalDelay' in conn else "0",
                        conn['gtfs:route'].replace('http://irail.be/vehicle/', ''),
                        conn['gtfs:trip'].replace('http://irail.be/trips/', ''),
                        conn['@id'].replace('http://irail.be/connections/', '')])
                if len(result) > batch_size:
                    return_value = result
                    result = []
{'fraction_non_alphanumeric': 0.07129053038436761, 'fraction_numerical': 0.01188175506406127, 'mean_word_length': 1.041794534560865, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 40, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1627867', 'n_tokens_mistral': 3442, 'n_tokens_neox': 2995, 'n_words': 842}