Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

# encoding: utf-8 

from __future__ import print_function, division, absolute_import 

 

import os 

import sys 

 

from sqlalchemy import create_engine, MetaData, Table 

from sqlalchemy.engine import reflection 

from sqlalchemy import exc 

from sqlalchemy.orm import sessionmaker 

 

from .db_objects import Base, declarative_base 

from .errors import InvalidOperationError 

from .utils import hexdigest, print_table 

from .logger import logger 

 

 

def connect_to_db(db_config, *, verbose=False): 

 

if db_config.connection_string.startswith("sqlite"): 

from sqlite3 import dbapi2 as sqlite 

engine = create_engine(db_config.connection_string, module=sqlite, echo=verbose) 

else: 

engine = create_engine(db_config.connection_string, echo=verbose) 

try: 

engine.connect() 

except exc.OperationalError as e: 

raise InvalidOperationError("could not connect to {}. Error is {}".format( 

db_config.connection_string, e)) from None 

l = logger() 

l.info("connected to db {}".format(db_config.connection_string)) 

return engine 

 

 

def setup_db(db_config, *, verbose=False): 

"""creates tables in database""" 

if check_if_tables_exist(db_config): 

raise InvalidOperationError("use setup_fresh_db, the tables in {!r} already exist".format( 

db_config.connection_string)) 

engine = connect_to_db(db_config, verbose=verbose) 

Base.metadata.create_all(engine) 

logger().info("created all tables of db {}".format(db_config.connection_string)) 

return engine 

 

 

def check_if_tables_exist(db_config, *, verbose=False): 

engine = connect_to_db(db_config, verbose=verbose) 

declared_names = Base.metadata.tables.keys() 

existing_names = reflection.Inspector.from_engine(engine).get_table_names() 

return bool(set(declared_names) & set(existing_names)) 

 

 

def setup_fresh_db(db_config, *, verbose=False): 

"""creates tables in database, deletes already existing data if present""" 

if not check_if_tables_exist(db_config): 

raise InvalidOperationError("use setup_db, the tables in {!r} do not exist".format( 

db_config.connection_string)) 

engine = connect_to_db(db_config, verbose=verbose) 

Base.metadata.drop_all(engine) 

Base.metadata.create_all(engine) 

logger().info("droped and created all tables of db {}".format(db_config.connection_string)) 

return engine 

 

 

def copy_db(db_config_source, db_config_destination, *, delete_existing, copy_signals, verbose=False): 

logger().info("copy tables {} -> {}".format(db_config_source.connection_string, 

db_config_destination.connection_string)) 

engine_source = connect_to_db(db_config_source, verbose=verbose) 

engine_destination = connect_to_db(db_config_destination, verbose=verbose) 

 

table_names_source = reflection.Inspector.from_engine(engine_source).get_table_names() 

existing_names_destination = reflection.Inspector.from_engine( 

engine_destination).get_table_names() 

 

if not delete_existing: 

if set(table_names_source) & set(existing_names_destination): 

raise InvalidOperationError("can not copy {} -> {}, some tables already exist on " 

"target db".format(db_config_source.connection_string, 

db_config_destination.connection_string)) 

 

for table_name in table_names_source: 

# TODO: import Dbo from db_objects.py and use the configure table name: 

# TODO: write test 

if table_name == "signals" and not copy_signals: 

continue 

_copy_table(engine_source, engine_destination, table_name, delete_existing, verbose) 

 

 

def _copy_table(engine_source, engine_destination, table_name, delete_existing, verbose): 

source = sessionmaker(bind=engine_source)() 

destination = sessionmaker(bind=engine_destination)() 

 

source_meta = MetaData(bind=engine_source) 

 

logger().info("copy schema of table {}".format(table_name)) 

 

table = Table(table_name, source_meta, autoload=True) 

if delete_existing: 

table.metadata.drop_all(engine_destination) 

table.metadata.create_all(engine_destination) 

 

Base = declarative_base() 

 

class NewRecord(Base): 

__table__ = table 

 

columns = table.columns.keys() 

logger().info("copy rows of table {}".format(table_name)) 

for record in source.query(table).all(): 

data = dict( 

[(str(column), getattr(record, column)) for column in columns] 

) 

destination.merge(NewRecord(**data)) 

 

logger().info("commit changes for table {}".format(table_name)) 

destination.commit() 

 

 

def _dump_table(engine, source, source_meta, table_name, indent="", file=sys.stdout): 

table = Table(table_name, source_meta, autoload=True) 

columns = table.columns.keys() 

rows = [] 

for record in source.query(table).all(): 

row = [] 

for column in columns: 

data = getattr(record, column) 

if isinstance(data, bytes): 

data = hexdigest(data) 

row.append(data) 

rows.append(row) 

print_table(columns, rows, indent=indent, file=file) 

 

 

def dump_db(db_config, file=sys.stdout): 

engine = connect_to_db(db_config) 

source = sessionmaker(bind=engine)() 

source_meta = MetaData(bind=engine) 

table_names = reflection.Inspector.from_engine(engine).get_table_names() 

for table_name in table_names: 

print("table {}:".format(table_name), file=file) 

print(file=file) 

_dump_table(engine, source, source_meta, table_name, indent=" ", file=file) 

print(file=file)