Completed
Push — master ( f2b7fb...01f125 )
by Roy
27s
created

_connect_elasticsearch()   B

Complexity

Conditions 7

Size

Total Lines 20

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 7
c 1
b 0
f 0
dl 0
loc 20
rs 7.3333
1
#!/usr/bin/env python
2
# -*- encoding: utf-8 -*-
3
# vim: set et sw=4 ts=4 sts=4 ff=unix fenc=utf8:
4
# Author: Binux<[email protected]>
5
#         http://binux.me
6
# Created on 2014-10-08 15:04:08
7
8
from six.moves.urllib.parse import urlparse, parse_qs
9
10
11
def connect_database(url):
12
    """
13
    create database object by url
14
15
    mysql:
16
        mysql+type://user:passwd@host:port/database
17
    sqlite:
18
        # relative path
19
        sqlite+type:///path/to/database.db
20
        # absolute path
21
        sqlite+type:////path/to/database.db
22
        # memory database
23
        sqlite+type://
24
    mongodb:
25
        mongodb+type://[username:password@]host1[:port1][,host2[:port2],...[,hostN[:portN]]][/[database][?options]]
26
        more: http://docs.mongodb.org/manual/reference/connection-string/
27
    sqlalchemy:
28
        sqlalchemy+postgresql+type://user:passwd@host:port/database
29
        sqlalchemy+mysql+mysqlconnector+type://user:passwd@host:port/database
30
        more: http://docs.sqlalchemy.org/en/rel_0_9/core/engines.html
31
    redis:
32
        redis+taskdb://host:port/db
33
    elasticsearch:
34
        elasticsearch+type://host:port/?index=pyspider
35
    local:
36
        local+projectdb://filepath,filepath
37
38
    type:
39
        taskdb
40
        projectdb
41
        resultdb
42
43
    """
44
    db = _connect_database(url)
45
    db.copy = lambda: _connect_database(url)
46
    return db
47
48
49
def _connect_database(url):  # NOQA
50
    parsed = urlparse(url)
51
52
    scheme = parsed.scheme.split('+')
53
    if len(scheme) == 1:
54
        raise Exception('wrong scheme format: %s' % parsed.scheme)
55
    else:
56
        engine, dbtype = scheme[0], scheme[-1]
57
        other_scheme = "+".join(scheme[1:-1])
58
59
    if dbtype not in ('taskdb', 'projectdb', 'resultdb'):
60
        raise LookupError('unknown database type: %s, '
61
                          'type should be one of ["taskdb", "projectdb", "resultdb"]', dbtype)
62
63
    if engine == 'mysql':
64
        return _connect_mysql(parsed,dbtype)
65
66
    elif engine == 'sqlite':
67
        return _connect_sqlite(parsed,dbtype)
68
    elif engine == 'mongodb':
69
        return _connect_mongodb(parsed,dbtype,url)
70
71
    elif engine == 'sqlalchemy':
72
        return _connect_sqlalchemy(parsed, dbtype, url, other_scheme)
73
74
75
    elif engine == 'redis':
76
        if dbtype == 'taskdb':
77
            from .redis.taskdb import TaskDB
78
            return TaskDB(parsed.hostname, parsed.port,
79
                          int(parsed.path.strip('/') or 0))
80
        else:
81
            raise LookupError('not supported dbtype: %s', dbtype)
82
    elif engine == 'local':
83
        scripts = url.split('//', 1)[1].split(',')
84
        if dbtype == 'projectdb':
85
            from .local.projectdb import ProjectDB
86
            return ProjectDB(scripts)
87
        else:
88
            raise LookupError('not supported dbtype: %s', dbtype)
89
    elif engine == 'elasticsearch' or engine == 'es':
90
        return _connect_elasticsearch(parsed, dbtype)
91
92
    else:
93
        raise Exception('unknown engine: %s' % engine)
94
95
96
def _connect_mysql(parsed,dbtype):
97
    parames = {}
98
    if parsed.username:
99
        parames['user'] = parsed.username
100
    if parsed.password:
101
        parames['passwd'] = parsed.password
102
    if parsed.hostname:
103
        parames['host'] = parsed.hostname
104
    if parsed.port:
105
        parames['port'] = parsed.port
106
    if parsed.path.strip('/'):
107
        parames['database'] = parsed.path.strip('/')
108
109
    if dbtype == 'taskdb':
110
        from .mysql.taskdb import TaskDB
111
        return TaskDB(**parames)
112
    elif dbtype == 'projectdb':
113
        from .mysql.projectdb import ProjectDB
114
        return ProjectDB(**parames)
115
    elif dbtype == 'resultdb':
116
        from .mysql.resultdb import ResultDB
117
        return ResultDB(**parames)
118
    else:
119
        raise LookupError
120
121
122
def _connect_sqlite(parsed,dbtype):
123
    if parsed.path.startswith('//'):
124
        path = '/' + parsed.path.strip('/')
125
    elif parsed.path.startswith('/'):
126
        path = './' + parsed.path.strip('/')
127
    elif not parsed.path:
128
        path = ':memory:'
129
    else:
130
        raise Exception('error path: %s' % parsed.path)
131
132
    if dbtype == 'taskdb':
133
        from .sqlite.taskdb import TaskDB
134
        return TaskDB(path)
135
    elif dbtype == 'projectdb':
136
        from .sqlite.projectdb import ProjectDB
137
        return ProjectDB(path)
138
    elif dbtype == 'resultdb':
139
        from .sqlite.resultdb import ResultDB
140
        return ResultDB(path)
141
    else:
142
        raise LookupError
143
144
145
def _connect_mongodb(parsed,dbtype,url):
146
    url = url.replace(parsed.scheme, 'mongodb')
147
    parames = {}
148
    if parsed.path.strip('/'):
149
        parames['database'] = parsed.path.strip('/')
150
151
    if dbtype == 'taskdb':
152
        from .mongodb.taskdb import TaskDB
153
        return TaskDB(url, **parames)
154
    elif dbtype == 'projectdb':
155
        from .mongodb.projectdb import ProjectDB
156
        return ProjectDB(url, **parames)
157
    elif dbtype == 'resultdb':
158
        from .mongodb.resultdb import ResultDB
159
        return ResultDB(url, **parames)
160
    else:
161
        raise LookupError
162
163
164
def _connect_sqlalchemy(parsed, dbtype,url, other_scheme):
165
    if not other_scheme:
166
        raise Exception('wrong scheme format: %s' % parsed.scheme)
167
    url = url.replace(parsed.scheme, other_scheme)
168
    if dbtype == 'taskdb':
169
        from .sqlalchemy.taskdb import TaskDB
170
        return TaskDB(url)
171
    elif dbtype == 'projectdb':
172
        from .sqlalchemy.projectdb import ProjectDB
173
        return ProjectDB(url)
174
    elif dbtype == 'resultdb':
175
        from .sqlalchemy.resultdb import ResultDB
176
        return ResultDB(url)
177
    else:
178
        raise LookupError
179
180
181
def _connect_elasticsearch(parsed, dbtype):
182
    # in python 2.6 url like "http://host/?query", query will not been splitted
183
    if parsed.path.startswith('/?'):
184
        index = parse_qs(parsed.path[2:])
185
    else:
186
        index = parse_qs(parsed.query)
187
    if 'index' in index and index['index']:
188
        index = index['index'][0]
189
    else:
190
        index = 'pyspider'
191
192
    if dbtype == 'projectdb':
193
        from .elasticsearch.projectdb import ProjectDB
194
        return ProjectDB([parsed.netloc], index=index)
195
    elif dbtype == 'resultdb':
196
        from .elasticsearch.resultdb import ResultDB
197
        return ResultDB([parsed.netloc], index=index)
198
    elif dbtype == 'taskdb':
199
        from .elasticsearch.taskdb import TaskDB
200
        return TaskDB([parsed.netloc], index=index)
201