ozbot/init_db.py at master · drJabber/ozbot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
#!flask/bin/python3

import sys
#from alembic import op,context
import sqlalchemy as sa
import geoalchemy2 as ga
import pandas as pd
import numpy as np
import csv


def run_migrations():
    metadata=sa.MetaData()
    table=create_stations_table(metadata)

    engine=sa.create_engine('postgresql://ozbot@localhost/ozbot')
#    connection = engine.connect().execution_options(schema_translate_map={None:"rzd"})
    connection = engine.connect()
    engine.execute('CREATE SCHEMA IF NOT EXISTS rzd')
    metadata.create_all(engine)
    populate(engine)
    connection.close()


def column_could_be_renamed_esr(column):
    if column in ('esr','express'):
        return column
    else:
        return column+'_esr'

def column_could_be_renamed_express(column):
    if column in ('esr','express'):
        return column
    else:
        return column+'_express'

def column_could_be_renamed_osm(column):
    if column in ('esr','express'):
        return column
    else:
        return column+'_osm'

def populate(engine):
    df=pd.read_csv('./oz/data/esr.csv',delimiter=';', dtype={'esr':object,'express':object, 'dup_esr':object })
    df.set_index([df['esr'],df['express']],inplace=True)
    df = df.rename(columns=column_could_be_renamed_esr)

    dfe=pd.read_csv('./oz/data/express.csv',delimiter=';', dtype={'esr':object,'express':object})
    dfe.set_index([dfe['esr'],dfe['express']],inplace=True)
    dfe = dfe.rename(columns=column_could_be_renamed_express)

    dfo=pd.read_csv('./oz/data/osm2esr.csv',delimiter=';', dtype={'esr':object})
    dfo.set_index([dfo['esr']],inplace=True)
    dfo = dfo.rename(columns=column_could_be_renamed_osm)

    dfm=pd.merge(df,dfe,on=['esr','express'],suffixes=['_esr','_express'],how='outer')
    dfm=pd.merge(dfm,dfo,on=['esr'],suffixes=['_esr','_osm'],how='outer')
    dfm.drop_duplicates(['esr','express'],keep='first',inplace=True)

    dfm.drop('user_osm',axis=1,inplace=True)


    engine.execute('delete from rzd.stations')
    dfm.to_sql('stations',engine,schema='rzd',if_exists='append',index=False)

    engine.execute('update rzd.stations set location=ST_SetSRID(ST_MakePoint(lon_osm, lat_osm),4326)::geometry where lon_osm is not null and lat_osm is not null')

def create_stations_table(metadata):
    ### commands auto generated by Alembic - please adjust! ###
    return sa.Table('stations',metadata,
    sa.Column('id', sa.Integer(), autoincrement=True,primary_key=True,nullable=False),
    sa.Column('division_esr', sa.String(length=200), nullable=True),
    sa.Column('esr', sa.String(length=12), nullable=True),
    sa.Column('country_esr', sa.String(length=200), nullable=True),
    sa.Column('region_esr', sa.String(length=200), nullable=True),
    sa.Column('express', sa.String(length=12), nullable=True),
    sa.Column('dup_esr_esr', sa.String(length=12), nullable=True),
    sa.Column('source_esr', sa.String(length=200), nullable=True),
    sa.Column('iso3166_esr', sa.String(length=100), nullable=True),
    sa.Column('railway_esr', sa.String(length=100), nullable=True),
    sa.Column('type_esr', sa.String(length=60), nullable=True),
    sa.Column('name_esr', sa.String(length=200), nullable=True),
    sa.Column('alias_express', sa.String(length=400), nullable=True),
    sa.Column('railway_express', sa.String(length=100), nullable=True),
    sa.Column('name_express', sa.String(length=400), nullable=True),
    sa.Column('status_osm', sa.Integer(), nullable=True),
    sa.Column('type_osm', sa.Integer(), nullable=True),
    sa.Column('osm_id_osm', sa.String(12), nullable=True),
    sa.Column('lat_osm', sa.Float(), nullable=True),
    sa.Column('lon_osm', sa.Float(), nullable=True),
    sa.Column('name_osm', sa.String(length=400), nullable=True),
    sa.Column('alt_name_osm', sa.String(length=200), nullable=True),
    sa.Column('old_name_osm', sa.String(length=200), nullable=True),
    sa.Column('official_name_osm', sa.String(length=400), nullable=True),
    sa.Column('railway_osm', sa.String(length=200), nullable=True),
    sa.Column('location',ga.Geography(geometry_type='POINT',srid=4326),nullable=True),
    sa.Index('stations_locations_idx','location',postgresql_using='gist'),
    schema='rzd'
    )


if __name__ == '__main__':
#    sys.exit(run_dump())
    sys.exit(run_migrations())