hdx-python-database 1.3.4__py3-none-any.whl → 1.3.6__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- hdx/database/__init__.py +2 -321
- hdx/database/_version.py +9 -4
- hdx/database/database.py +345 -0
- hdx/database/views.py +8 -17
- {hdx_python_database-1.3.4.dist-info → hdx_python_database-1.3.6.dist-info}/METADATA +4 -2
- hdx_python_database-1.3.6.dist-info/RECORD +13 -0
- {hdx_python_database-1.3.4.dist-info → hdx_python_database-1.3.6.dist-info}/WHEEL +1 -1
- hdx_python_database-1.3.4.dist-info/RECORD +0 -12
- {hdx_python_database-1.3.4.dist-info → hdx_python_database-1.3.6.dist-info}/licenses/LICENSE +0 -0
hdx/database/__init__.py
CHANGED
|
@@ -1,322 +1,3 @@
|
|
|
1
|
-
"""Database utilities"""
|
|
2
|
-
|
|
3
|
-
import logging
|
|
4
|
-
from typing import Any, Dict, List, Optional, Tuple, Type, Union
|
|
5
|
-
|
|
6
|
-
from sqlalchemy import Engine, TableClause, create_engine
|
|
7
|
-
from sqlalchemy.exc import SQLAlchemyError
|
|
8
|
-
from sqlalchemy.ext.automap import automap_base
|
|
9
|
-
from sqlalchemy.orm import DeclarativeBase, Session
|
|
10
|
-
from sqlalchemy.pool import NullPool
|
|
11
|
-
from sqlalchemy.sql.ddl import CreateSchema, DropSchema
|
|
12
|
-
|
|
13
1
|
from ._version import version as __version__ # noqa: F401
|
|
14
|
-
from .
|
|
15
|
-
from .
|
|
16
|
-
from .postgresql import restore_from_pgfile, wait_for_postgresql
|
|
17
|
-
from .views import view
|
|
18
|
-
from .with_timezone import Base as TZBase
|
|
19
|
-
|
|
20
|
-
logger = logging.getLogger(__name__)
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
class DatabaseError(Exception):
|
|
24
|
-
pass
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
def do_nothing_fn():
|
|
28
|
-
return None
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
class Database:
|
|
32
|
-
"""Database helper class to handle ssh tunnels, waiting for PostgreSQL to
|
|
33
|
-
be up etc. Can be used in a with statement returning a Database object that
|
|
34
|
-
if reflect is True will have a variable reflected_classes containing the
|
|
35
|
-
reflected classes. Either an SQLAlchemy engine, database uri or various
|
|
36
|
-
database parameters must be supplied.
|
|
37
|
-
|
|
38
|
-
db_has_tz which defaults to False indicates whether database datetime
|
|
39
|
-
columns have timezones. If not, conversion occurs between Python datetimes
|
|
40
|
-
with timezones to timezoneless database columns (but not when using
|
|
41
|
-
reflection). If table_base is supplied, db_has_tz is ignored.
|
|
42
|
-
|
|
43
|
-
The database can be restored from a given pg_restore file supplied in the
|
|
44
|
-
parameter pg_restore_file.
|
|
45
|
-
|
|
46
|
-
There is an option to wipe and create an empty schema in the database by
|
|
47
|
-
setting recreate_schema to True and setting a schema_name ("public" is the
|
|
48
|
-
default).
|
|
49
|
-
|
|
50
|
-
If a prepare function is supplied in prepare_fn, it will be executed before
|
|
51
|
-
Base.metadata.create_all and the results of it returned in instance variable
|
|
52
|
-
prepare_results.
|
|
53
|
-
|
|
54
|
-
Args:
|
|
55
|
-
engine (Optional[Engine]): SQLAlchemy engine to use.
|
|
56
|
-
db_uri (Optional[str]): Connection URI.
|
|
57
|
-
database (Optional[str]): Database name
|
|
58
|
-
host (Optional[str]): Host where database is located
|
|
59
|
-
port (Union[int, str, None]): Database port
|
|
60
|
-
username (Optional[str]): Username to log into database
|
|
61
|
-
password (Optional[str]): Password to log into database
|
|
62
|
-
dialect (str): Database dialect. Defaults to "postgresql".
|
|
63
|
-
driver (Optional[str]): Database driver. Defaults to None (psycopg if postgresql or None)
|
|
64
|
-
db_has_tz (bool): True if db datetime columns have timezone. Defaults to False.
|
|
65
|
-
table_base (Optional[Type[DeclarativeBase]]): Override table base. Defaults to None.
|
|
66
|
-
reflect (bool): Whether to reflect existing tables. Defaults to False.
|
|
67
|
-
**kwargs: See below
|
|
68
|
-
pg_restore_file (str): Restore database from pg_restore file
|
|
69
|
-
recreate_schema (bool): Whether to recreate schema
|
|
70
|
-
schema_name (str): Database schema name. Defaults to "public".
|
|
71
|
-
prepare_fn (Callable[[], None]]): Function to call before Base.metadata.create_all.
|
|
72
|
-
ssh_host (str): SSH host (the server to connect to)
|
|
73
|
-
ssh_port (int): SSH port. Defaults to 22.
|
|
74
|
-
ssh_username (str): SSH username
|
|
75
|
-
ssh_password (str): SSH password
|
|
76
|
-
ssh_private_key: Location of SSH private key (instead of ssh_password)
|
|
77
|
-
For more advanced usage, see SSHTunnelForwarder documentation.
|
|
78
|
-
|
|
79
|
-
"""
|
|
80
|
-
|
|
81
|
-
def __init__(
|
|
82
|
-
self,
|
|
83
|
-
engine: Optional[Engine] = None,
|
|
84
|
-
db_uri: Optional[str] = None,
|
|
85
|
-
database: Optional[str] = None,
|
|
86
|
-
host: Optional[str] = None,
|
|
87
|
-
port: Union[int, str, None] = None,
|
|
88
|
-
username: Optional[str] = None,
|
|
89
|
-
password: Optional[str] = None,
|
|
90
|
-
dialect: str = "postgresql",
|
|
91
|
-
driver: Optional[str] = None,
|
|
92
|
-
db_has_tz: bool = False,
|
|
93
|
-
table_base: Optional[Type[DeclarativeBase]] = None,
|
|
94
|
-
reflect: bool = False,
|
|
95
|
-
**kwargs: Any,
|
|
96
|
-
) -> None:
|
|
97
|
-
if port is not None:
|
|
98
|
-
port = int(port)
|
|
99
|
-
schema_name = None
|
|
100
|
-
if len(kwargs) == 0:
|
|
101
|
-
pg_restore_file = None
|
|
102
|
-
recreate_schema = False
|
|
103
|
-
prepare_fn = do_nothing_fn
|
|
104
|
-
else:
|
|
105
|
-
pg_restore_file = kwargs.pop("pg_restore_file", None)
|
|
106
|
-
recreate_schema = kwargs.pop("recreate_schema", False)
|
|
107
|
-
schema_name = kwargs.pop("schema", "public")
|
|
108
|
-
prepare_fn = kwargs.pop("prepare_fn", do_nothing_fn)
|
|
109
|
-
if len(kwargs) != 0:
|
|
110
|
-
try:
|
|
111
|
-
import sshtunnel
|
|
112
|
-
except ImportError:
|
|
113
|
-
# dependency missing, log an error
|
|
114
|
-
logger.error(
|
|
115
|
-
"sshtunnel not found! Please install hdx-python-database[sshtunnel] to enable."
|
|
116
|
-
)
|
|
117
|
-
raise
|
|
118
|
-
ssh_host = kwargs["ssh_host"]
|
|
119
|
-
del kwargs["ssh_host"]
|
|
120
|
-
ssh_port = kwargs.get("ssh_port")
|
|
121
|
-
if ssh_port is not None:
|
|
122
|
-
ssh_port = int(ssh_port)
|
|
123
|
-
del kwargs["ssh_port"]
|
|
124
|
-
else:
|
|
125
|
-
ssh_port = 22
|
|
126
|
-
self.server = sshtunnel.SSHTunnelForwarder(
|
|
127
|
-
(ssh_host, ssh_port),
|
|
128
|
-
remote_bind_address=(host, port),
|
|
129
|
-
**kwargs,
|
|
130
|
-
)
|
|
131
|
-
self.server.start()
|
|
132
|
-
host = self.server.local_bind_host
|
|
133
|
-
port = self.server.local_bind_port
|
|
134
|
-
else:
|
|
135
|
-
self.server = None
|
|
136
|
-
if not engine and not db_uri:
|
|
137
|
-
db_uri = get_connection_uri(
|
|
138
|
-
database,
|
|
139
|
-
host,
|
|
140
|
-
port,
|
|
141
|
-
username,
|
|
142
|
-
password,
|
|
143
|
-
dialect=dialect,
|
|
144
|
-
driver=driver,
|
|
145
|
-
)
|
|
146
|
-
if not db_uri:
|
|
147
|
-
raise DatabaseError(
|
|
148
|
-
"No engine, database uri or connection parameters provided!"
|
|
149
|
-
)
|
|
150
|
-
if not engine and dialect == "postgresql":
|
|
151
|
-
wait_for_postgresql(db_uri)
|
|
152
|
-
if pg_restore_file and db_uri:
|
|
153
|
-
restore_from_pgfile(db_uri, pg_restore_file)
|
|
154
|
-
if not table_base:
|
|
155
|
-
if db_has_tz:
|
|
156
|
-
table_base = TZBase
|
|
157
|
-
else:
|
|
158
|
-
table_base = NoTZBase
|
|
159
|
-
if not engine:
|
|
160
|
-
engine = create_engine(db_uri, poolclass=NullPool, echo=False)
|
|
161
|
-
self.engine: Engine = engine
|
|
162
|
-
if recreate_schema:
|
|
163
|
-
self.recreate_schema(engine, schema_name)
|
|
164
|
-
self.prepare_results = prepare_fn()
|
|
165
|
-
self.session, self.base = self.create_session(
|
|
166
|
-
engine,
|
|
167
|
-
table_base=table_base,
|
|
168
|
-
reflect=reflect,
|
|
169
|
-
)
|
|
170
|
-
if reflect:
|
|
171
|
-
self.reflected_classes = self.base.classes
|
|
172
|
-
else:
|
|
173
|
-
self.reflected_classes = None
|
|
174
|
-
|
|
175
|
-
def cleanup(self) -> None:
|
|
176
|
-
"""Cleanup SQLAlchemy.
|
|
177
|
-
|
|
178
|
-
Returns:
|
|
179
|
-
sqlalchemy.Engine: SQLAlchemy engine
|
|
180
|
-
"""
|
|
181
|
-
self.session.close()
|
|
182
|
-
self.engine.dispose()
|
|
183
|
-
if self.server is not None:
|
|
184
|
-
self.server.stop()
|
|
185
|
-
|
|
186
|
-
def __enter__(self) -> "Database":
|
|
187
|
-
return self
|
|
188
|
-
|
|
189
|
-
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
|
190
|
-
self.cleanup()
|
|
191
|
-
|
|
192
|
-
def drop_all(self) -> None:
|
|
193
|
-
"""Drop all tables.
|
|
194
|
-
|
|
195
|
-
Returns:
|
|
196
|
-
None
|
|
197
|
-
"""
|
|
198
|
-
self.base.metadata.drop_all(self.engine)
|
|
199
|
-
|
|
200
|
-
def get_engine(self) -> Engine:
|
|
201
|
-
"""Returns SQLAlchemy engine.
|
|
202
|
-
|
|
203
|
-
Returns:
|
|
204
|
-
sqlalchemy.Engine: SQLAlchemy engine
|
|
205
|
-
"""
|
|
206
|
-
return self.engine
|
|
207
|
-
|
|
208
|
-
def get_session(self) -> Session:
|
|
209
|
-
"""Returns SQLAlchemy session.
|
|
210
|
-
|
|
211
|
-
Returns:
|
|
212
|
-
sqlalchemy.orm.Session: SQLAlchemy session
|
|
213
|
-
"""
|
|
214
|
-
return self.session
|
|
215
|
-
|
|
216
|
-
def get_reflected_classes(self) -> Any:
|
|
217
|
-
"""Gets reflected classes.
|
|
218
|
-
|
|
219
|
-
Returns:
|
|
220
|
-
Any: Reflected classes
|
|
221
|
-
"""
|
|
222
|
-
return self.reflected_classes
|
|
223
|
-
|
|
224
|
-
def get_prepare_results(self) -> Any:
|
|
225
|
-
"""Returns results from prepare function.
|
|
226
|
-
|
|
227
|
-
Returns:
|
|
228
|
-
Any: Results from prepare function
|
|
229
|
-
"""
|
|
230
|
-
return self.prepare_results
|
|
231
|
-
|
|
232
|
-
@staticmethod
|
|
233
|
-
def create_session(
|
|
234
|
-
engine: Optional[Engine] = None,
|
|
235
|
-
db_uri: Optional[str] = None,
|
|
236
|
-
table_base: Type[DeclarativeBase] = NoTZBase,
|
|
237
|
-
reflect: bool = False,
|
|
238
|
-
) -> Tuple[Session, Any]:
|
|
239
|
-
"""Creates SQLAlchemy session given SQLAlchemy engine or database uri
|
|
240
|
-
(one of which must be supplied). Tables must inherit from Base in
|
|
241
|
-
hdx.utilities.database unless base is defined. If reflect is True,
|
|
242
|
-
classes will be reflected from an existing database and the reflected
|
|
243
|
-
classes will be returned. Note that type annotation maps don't work
|
|
244
|
-
with reflection.
|
|
245
|
-
|
|
246
|
-
Args:
|
|
247
|
-
engine (Optional[Engine]): SQLAlchemy engine to use. Defaults to None (create from db_uri).
|
|
248
|
-
db_uri (Optional[str]): Connection URI. Defaults to None (use engine).
|
|
249
|
-
table_base (Type[DeclarativeBase]): Base database table class. Defaults to NoTZBase.
|
|
250
|
-
reflect (bool): Whether to reflect existing tables. Defaults to False.
|
|
251
|
-
|
|
252
|
-
Returns:
|
|
253
|
-
Tuple[Session, Any]: (SQLAlchemy session, base)
|
|
254
|
-
"""
|
|
255
|
-
if not engine:
|
|
256
|
-
if db_uri is None:
|
|
257
|
-
raise DatabaseError("No engine or database uri provided!")
|
|
258
|
-
engine = create_engine(db_uri, poolclass=NullPool, echo=False)
|
|
259
|
-
if reflect:
|
|
260
|
-
Base = automap_base(declarative_base=table_base)
|
|
261
|
-
Base.prepare(
|
|
262
|
-
autoload_with=engine, reflection_options={"views": True}
|
|
263
|
-
)
|
|
264
|
-
table_base = Base
|
|
265
|
-
else:
|
|
266
|
-
table_base.metadata.create_all(engine)
|
|
267
|
-
return Session(engine), table_base
|
|
268
|
-
|
|
269
|
-
@staticmethod
|
|
270
|
-
def recreate_schema(engine: Engine, schema_name: str = "public") -> bool:
|
|
271
|
-
"""Wipe and create empty schema in database using SQLAlchemy.
|
|
272
|
-
|
|
273
|
-
Args:
|
|
274
|
-
engine (Engine): SQLAlchemy engine to use.
|
|
275
|
-
schema_name (str): Schema name. Defaults to "public".
|
|
276
|
-
|
|
277
|
-
Returns:
|
|
278
|
-
bool: True if all successful, False if not
|
|
279
|
-
"""
|
|
280
|
-
# Wipe and create an empty schema
|
|
281
|
-
try:
|
|
282
|
-
with engine.connect() as connection:
|
|
283
|
-
connection.execute(
|
|
284
|
-
DropSchema(schema_name, cascade=True, if_exists=True)
|
|
285
|
-
)
|
|
286
|
-
connection.commit()
|
|
287
|
-
connection.execute(
|
|
288
|
-
CreateSchema(schema_name, if_not_exists=True)
|
|
289
|
-
)
|
|
290
|
-
connection.commit()
|
|
291
|
-
return True
|
|
292
|
-
except SQLAlchemyError:
|
|
293
|
-
return False
|
|
294
|
-
|
|
295
|
-
@staticmethod
|
|
296
|
-
def prepare_view(view_params: Dict) -> TableClause:
|
|
297
|
-
"""Prepare SQLAlchemy view from dictionary with keys: name, metadata and
|
|
298
|
-
selectable. Must be run before Base.metadata.create_all.
|
|
299
|
-
|
|
300
|
-
Args:
|
|
301
|
-
view_params (Dict): Dictionary with keys name, metadata, selectable
|
|
302
|
-
|
|
303
|
-
Returns:
|
|
304
|
-
TableClause: SQLAlchemy View
|
|
305
|
-
"""
|
|
306
|
-
return view(**view_params)
|
|
307
|
-
|
|
308
|
-
@classmethod
|
|
309
|
-
def prepare_views(cls, view_params_list: List[Dict]) -> List[TableClause]:
|
|
310
|
-
"""Prepare SQLAlchemy views from a list of dictionaries with keys: name,
|
|
311
|
-
metadata and selectable. Must be run before Base.metadata.create_all.
|
|
312
|
-
|
|
313
|
-
Args:
|
|
314
|
-
view_params_list (List[Dict]): List of dictionaries with view parameters
|
|
315
|
-
|
|
316
|
-
Returns:
|
|
317
|
-
List[TableClause]: SQLAlchemy Views
|
|
318
|
-
"""
|
|
319
|
-
results = []
|
|
320
|
-
for view_params in view_params_list:
|
|
321
|
-
results.append(cls.prepare_view(view_params))
|
|
322
|
-
return results
|
|
2
|
+
from .database import Database as Database
|
|
3
|
+
from .database import DatabaseError as DatabaseError
|
hdx/database/_version.py
CHANGED
|
@@ -1,8 +1,13 @@
|
|
|
1
|
-
# file generated by
|
|
1
|
+
# file generated by setuptools-scm
|
|
2
2
|
# don't change, don't track in version control
|
|
3
|
+
|
|
4
|
+
__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
|
|
5
|
+
|
|
3
6
|
TYPE_CHECKING = False
|
|
4
7
|
if TYPE_CHECKING:
|
|
5
|
-
from typing import Tuple
|
|
8
|
+
from typing import Tuple
|
|
9
|
+
from typing import Union
|
|
10
|
+
|
|
6
11
|
VERSION_TUPLE = Tuple[Union[int, str], ...]
|
|
7
12
|
else:
|
|
8
13
|
VERSION_TUPLE = object
|
|
@@ -12,5 +17,5 @@ __version__: str
|
|
|
12
17
|
__version_tuple__: VERSION_TUPLE
|
|
13
18
|
version_tuple: VERSION_TUPLE
|
|
14
19
|
|
|
15
|
-
__version__ = version = '1.3.
|
|
16
|
-
__version_tuple__ = version_tuple = (1, 3,
|
|
20
|
+
__version__ = version = '1.3.6'
|
|
21
|
+
__version_tuple__ = version_tuple = (1, 3, 6)
|
hdx/database/database.py
ADDED
|
@@ -0,0 +1,345 @@
|
|
|
1
|
+
"""Database utilities"""
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
from typing import Any, Dict, List, Optional, Tuple, Type, Union
|
|
5
|
+
|
|
6
|
+
from sqlalchemy import Engine, TableClause, create_engine, insert
|
|
7
|
+
from sqlalchemy.exc import SQLAlchemyError
|
|
8
|
+
from sqlalchemy.ext.automap import automap_base
|
|
9
|
+
from sqlalchemy.orm import DeclarativeBase, Session
|
|
10
|
+
from sqlalchemy.pool import NullPool
|
|
11
|
+
from sqlalchemy.sql.ddl import CreateSchema, DropSchema
|
|
12
|
+
|
|
13
|
+
from .dburi import get_connection_uri
|
|
14
|
+
from .no_timezone import Base as NoTZBase
|
|
15
|
+
from .postgresql import restore_from_pgfile, wait_for_postgresql
|
|
16
|
+
from .views import view
|
|
17
|
+
from .with_timezone import Base as TZBase
|
|
18
|
+
|
|
19
|
+
logger = logging.getLogger(__name__)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class DatabaseError(Exception):
|
|
23
|
+
pass
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def do_nothing_fn():
|
|
27
|
+
return None
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class Database:
|
|
31
|
+
"""Database helper class to handle ssh tunnels, waiting for PostgreSQL to
|
|
32
|
+
be up etc. Can be used in a with statement returning a Database object that
|
|
33
|
+
if reflect is True will have a variable reflected_classes containing the
|
|
34
|
+
reflected classes. Either an SQLAlchemy engine, database uri or various
|
|
35
|
+
database parameters must be supplied.
|
|
36
|
+
|
|
37
|
+
db_has_tz which defaults to False indicates whether database datetime
|
|
38
|
+
columns have timezones. If not, conversion occurs between Python datetimes
|
|
39
|
+
with timezones to timezoneless database columns (but not when using
|
|
40
|
+
reflection). If table_base is supplied, db_has_tz is ignored.
|
|
41
|
+
|
|
42
|
+
The database can be restored from a given pg_restore file supplied in the
|
|
43
|
+
parameter pg_restore_file.
|
|
44
|
+
|
|
45
|
+
There is an option to wipe and create an empty schema in the database by
|
|
46
|
+
setting recreate_schema to True and setting a schema_name ("public" is the
|
|
47
|
+
default).
|
|
48
|
+
|
|
49
|
+
If a prepare function is supplied in prepare_fn, it will be executed before
|
|
50
|
+
Base.metadata.create_all and the results of it returned in instance variable
|
|
51
|
+
prepare_results.
|
|
52
|
+
|
|
53
|
+
Args:
|
|
54
|
+
engine (Optional[Engine]): SQLAlchemy engine to use.
|
|
55
|
+
db_uri (Optional[str]): Connection URI.
|
|
56
|
+
database (Optional[str]): Database name
|
|
57
|
+
host (Optional[str]): Host where database is located
|
|
58
|
+
port (Union[int, str, None]): Database port
|
|
59
|
+
username (Optional[str]): Username to log into database
|
|
60
|
+
password (Optional[str]): Password to log into database
|
|
61
|
+
dialect (str): Database dialect. Defaults to "postgresql".
|
|
62
|
+
driver (Optional[str]): Database driver. Defaults to None (psycopg if postgresql or None)
|
|
63
|
+
db_has_tz (bool): True if db datetime columns have timezone. Defaults to False.
|
|
64
|
+
table_base (Optional[Type[DeclarativeBase]]): Override table base. Defaults to None.
|
|
65
|
+
reflect (bool): Whether to reflect existing tables. Defaults to False.
|
|
66
|
+
**kwargs: See below
|
|
67
|
+
pg_restore_file (str): Restore database from pg_restore file
|
|
68
|
+
recreate_schema (bool): Whether to recreate schema
|
|
69
|
+
schema_name (str): Database schema name. Defaults to "public".
|
|
70
|
+
prepare_fn (Callable[[], None]]): Function to call before Base.metadata.create_all.
|
|
71
|
+
ssh_host (str): SSH host (the server to connect to)
|
|
72
|
+
ssh_port (int): SSH port. Defaults to 22.
|
|
73
|
+
ssh_username (str): SSH username
|
|
74
|
+
ssh_password (str): SSH password
|
|
75
|
+
ssh_private_key: Location of SSH private key (instead of ssh_password)
|
|
76
|
+
For more advanced usage, see SSHTunnelForwarder documentation.
|
|
77
|
+
|
|
78
|
+
"""
|
|
79
|
+
|
|
80
|
+
def __init__(
|
|
81
|
+
self,
|
|
82
|
+
engine: Optional[Engine] = None,
|
|
83
|
+
db_uri: Optional[str] = None,
|
|
84
|
+
database: Optional[str] = None,
|
|
85
|
+
host: Optional[str] = None,
|
|
86
|
+
port: Union[int, str, None] = None,
|
|
87
|
+
username: Optional[str] = None,
|
|
88
|
+
password: Optional[str] = None,
|
|
89
|
+
dialect: str = "postgresql",
|
|
90
|
+
driver: Optional[str] = None,
|
|
91
|
+
db_has_tz: bool = False,
|
|
92
|
+
table_base: Optional[Type[DeclarativeBase]] = None,
|
|
93
|
+
reflect: bool = False,
|
|
94
|
+
**kwargs: Any,
|
|
95
|
+
) -> None:
|
|
96
|
+
if port is not None:
|
|
97
|
+
port = int(port)
|
|
98
|
+
schema_name = None
|
|
99
|
+
if len(kwargs) == 0:
|
|
100
|
+
pg_restore_file = None
|
|
101
|
+
recreate_schema = False
|
|
102
|
+
prepare_fn = do_nothing_fn
|
|
103
|
+
else:
|
|
104
|
+
pg_restore_file = kwargs.pop("pg_restore_file", None)
|
|
105
|
+
recreate_schema = kwargs.pop("recreate_schema", False)
|
|
106
|
+
schema_name = kwargs.pop("schema", "public")
|
|
107
|
+
prepare_fn = kwargs.pop("prepare_fn", do_nothing_fn)
|
|
108
|
+
if len(kwargs) != 0:
|
|
109
|
+
try:
|
|
110
|
+
import sshtunnel
|
|
111
|
+
except ImportError:
|
|
112
|
+
# dependency missing, log an error
|
|
113
|
+
logger.error(
|
|
114
|
+
"sshtunnel not found! Please install hdx-python-database[sshtunnel] to enable."
|
|
115
|
+
)
|
|
116
|
+
raise
|
|
117
|
+
ssh_host = kwargs["ssh_host"]
|
|
118
|
+
del kwargs["ssh_host"]
|
|
119
|
+
ssh_port = kwargs.get("ssh_port")
|
|
120
|
+
if ssh_port is not None:
|
|
121
|
+
ssh_port = int(ssh_port)
|
|
122
|
+
del kwargs["ssh_port"]
|
|
123
|
+
else:
|
|
124
|
+
ssh_port = 22
|
|
125
|
+
self._server = sshtunnel.SSHTunnelForwarder(
|
|
126
|
+
(ssh_host, ssh_port),
|
|
127
|
+
remote_bind_address=(host, port),
|
|
128
|
+
**kwargs,
|
|
129
|
+
)
|
|
130
|
+
self._server.start()
|
|
131
|
+
host = self._server.local_bind_host
|
|
132
|
+
port = self._server.local_bind_port
|
|
133
|
+
else:
|
|
134
|
+
self._server = None
|
|
135
|
+
if not engine and not db_uri:
|
|
136
|
+
db_uri = get_connection_uri(
|
|
137
|
+
database,
|
|
138
|
+
host,
|
|
139
|
+
port,
|
|
140
|
+
username,
|
|
141
|
+
password,
|
|
142
|
+
dialect=dialect,
|
|
143
|
+
driver=driver,
|
|
144
|
+
)
|
|
145
|
+
if not db_uri:
|
|
146
|
+
raise DatabaseError(
|
|
147
|
+
"No engine, database uri or connection parameters provided!"
|
|
148
|
+
)
|
|
149
|
+
if not engine and dialect == "postgresql":
|
|
150
|
+
wait_for_postgresql(db_uri)
|
|
151
|
+
if pg_restore_file and db_uri:
|
|
152
|
+
restore_from_pgfile(db_uri, pg_restore_file)
|
|
153
|
+
if not table_base:
|
|
154
|
+
if db_has_tz:
|
|
155
|
+
table_base = TZBase
|
|
156
|
+
else:
|
|
157
|
+
table_base = NoTZBase
|
|
158
|
+
if not engine:
|
|
159
|
+
engine = create_engine(db_uri, poolclass=NullPool, echo=False)
|
|
160
|
+
self._engine: Engine = engine
|
|
161
|
+
if recreate_schema:
|
|
162
|
+
self.recreate_schema(engine, schema_name)
|
|
163
|
+
self._prepare_results = prepare_fn()
|
|
164
|
+
self._session, self._base = self.create_session(
|
|
165
|
+
engine,
|
|
166
|
+
table_base=table_base,
|
|
167
|
+
reflect=reflect,
|
|
168
|
+
)
|
|
169
|
+
if reflect:
|
|
170
|
+
self._reflected_classes = self._base.classes
|
|
171
|
+
else:
|
|
172
|
+
self._reflected_classes = None
|
|
173
|
+
|
|
174
|
+
def cleanup(self) -> None:
|
|
175
|
+
"""Cleanup SQLAlchemy.
|
|
176
|
+
|
|
177
|
+
Returns:
|
|
178
|
+
sqlalchemy.Engine: SQLAlchemy engine
|
|
179
|
+
"""
|
|
180
|
+
self._session.close()
|
|
181
|
+
self._engine.dispose()
|
|
182
|
+
if self._server is not None:
|
|
183
|
+
self._server.stop()
|
|
184
|
+
|
|
185
|
+
def __enter__(self) -> "Database":
|
|
186
|
+
return self
|
|
187
|
+
|
|
188
|
+
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
|
189
|
+
self.cleanup()
|
|
190
|
+
|
|
191
|
+
def drop_all(self) -> None:
|
|
192
|
+
"""Drop all tables.
|
|
193
|
+
|
|
194
|
+
Returns:
|
|
195
|
+
None
|
|
196
|
+
"""
|
|
197
|
+
self._base.metadata.drop_all(self._engine)
|
|
198
|
+
|
|
199
|
+
def get_engine(self) -> Engine:
|
|
200
|
+
"""Returns SQLAlchemy engine.
|
|
201
|
+
|
|
202
|
+
Returns:
|
|
203
|
+
sqlalchemy.Engine: SQLAlchemy engine
|
|
204
|
+
"""
|
|
205
|
+
return self._engine
|
|
206
|
+
|
|
207
|
+
def get_session(self) -> Session:
|
|
208
|
+
"""Returns SQLAlchemy session.
|
|
209
|
+
|
|
210
|
+
Returns:
|
|
211
|
+
sqlalchemy.orm.Session: SQLAlchemy session
|
|
212
|
+
"""
|
|
213
|
+
return self._session
|
|
214
|
+
|
|
215
|
+
def get_reflected_classes(self) -> Any:
|
|
216
|
+
"""Gets reflected classes.
|
|
217
|
+
|
|
218
|
+
Returns:
|
|
219
|
+
Any: Reflected classes
|
|
220
|
+
"""
|
|
221
|
+
return self._reflected_classes
|
|
222
|
+
|
|
223
|
+
def get_prepare_results(self) -> Any:
|
|
224
|
+
"""Returns results from prepare function.
|
|
225
|
+
|
|
226
|
+
Returns:
|
|
227
|
+
Any: Results from prepare function
|
|
228
|
+
"""
|
|
229
|
+
return self._prepare_results
|
|
230
|
+
|
|
231
|
+
def batch_populate(
|
|
232
|
+
self,
|
|
233
|
+
rows: List[Dict],
|
|
234
|
+
dbtable: Type[DeclarativeBase],
|
|
235
|
+
batch_size: int = 1000,
|
|
236
|
+
) -> None:
|
|
237
|
+
"""Batch populate database table.
|
|
238
|
+
|
|
239
|
+
Args:
|
|
240
|
+
rows (List[Dict]): List of rows
|
|
241
|
+
dbtable (Type[DeclarativeBase]): Database table
|
|
242
|
+
batch_size (int): Batch size. Defaults to 1000.
|
|
243
|
+
|
|
244
|
+
Returns:
|
|
245
|
+
Tuple[Session, Any]: (SQLAlchemy session, base)
|
|
246
|
+
"""
|
|
247
|
+
batches = range(len(rows) // batch_size + 1)
|
|
248
|
+
for batch in batches:
|
|
249
|
+
start_row = batch * batch_size
|
|
250
|
+
end_row = start_row + batch_size
|
|
251
|
+
batch_rows = rows[start_row:end_row]
|
|
252
|
+
self._session.execute(insert(dbtable), batch_rows)
|
|
253
|
+
self._session.commit()
|
|
254
|
+
|
|
255
|
+
@staticmethod
|
|
256
|
+
def create_session(
|
|
257
|
+
engine: Optional[Engine] = None,
|
|
258
|
+
db_uri: Optional[str] = None,
|
|
259
|
+
table_base: Type[DeclarativeBase] = NoTZBase,
|
|
260
|
+
reflect: bool = False,
|
|
261
|
+
) -> Tuple[Session, Any]:
|
|
262
|
+
"""Creates SQLAlchemy session given SQLAlchemy engine or database uri
|
|
263
|
+
(one of which must be supplied). Tables must inherit from Base in
|
|
264
|
+
hdx.utilities.database unless base is defined. If reflect is True,
|
|
265
|
+
classes will be reflected from an existing database and the reflected
|
|
266
|
+
classes will be returned. Note that type annotation maps don't work
|
|
267
|
+
with reflection.
|
|
268
|
+
|
|
269
|
+
Args:
|
|
270
|
+
engine (Optional[Engine]): SQLAlchemy engine to use. Defaults to None (create from db_uri).
|
|
271
|
+
db_uri (Optional[str]): Connection URI. Defaults to None (use engine).
|
|
272
|
+
table_base (Type[DeclarativeBase]): Base database table class. Defaults to NoTZBase.
|
|
273
|
+
reflect (bool): Whether to reflect existing tables. Defaults to False.
|
|
274
|
+
|
|
275
|
+
Returns:
|
|
276
|
+
Tuple[Session, Any]: (SQLAlchemy session, base)
|
|
277
|
+
"""
|
|
278
|
+
if not engine:
|
|
279
|
+
if db_uri is None:
|
|
280
|
+
raise DatabaseError("No engine or database uri provided!")
|
|
281
|
+
engine = create_engine(db_uri, poolclass=NullPool, echo=False)
|
|
282
|
+
if reflect:
|
|
283
|
+
Base = automap_base(declarative_base=table_base)
|
|
284
|
+
Base.prepare(
|
|
285
|
+
autoload_with=engine, reflection_options={"views": True}
|
|
286
|
+
)
|
|
287
|
+
table_base = Base
|
|
288
|
+
else:
|
|
289
|
+
table_base.metadata.create_all(engine)
|
|
290
|
+
return Session(engine), table_base
|
|
291
|
+
|
|
292
|
+
@staticmethod
|
|
293
|
+
def recreate_schema(engine: Engine, schema_name: str = "public") -> bool:
|
|
294
|
+
"""Wipe and create empty schema in database using SQLAlchemy.
|
|
295
|
+
|
|
296
|
+
Args:
|
|
297
|
+
engine (Engine): SQLAlchemy engine to use.
|
|
298
|
+
schema_name (str): Schema name. Defaults to "public".
|
|
299
|
+
|
|
300
|
+
Returns:
|
|
301
|
+
bool: True if all successful, False if not
|
|
302
|
+
"""
|
|
303
|
+
# Wipe and create an empty schema
|
|
304
|
+
try:
|
|
305
|
+
with engine.connect() as connection:
|
|
306
|
+
connection.execute(
|
|
307
|
+
DropSchema(schema_name, cascade=True, if_exists=True)
|
|
308
|
+
)
|
|
309
|
+
connection.commit()
|
|
310
|
+
connection.execute(
|
|
311
|
+
CreateSchema(schema_name, if_not_exists=True)
|
|
312
|
+
)
|
|
313
|
+
connection.commit()
|
|
314
|
+
return True
|
|
315
|
+
except SQLAlchemyError:
|
|
316
|
+
return False
|
|
317
|
+
|
|
318
|
+
@staticmethod
|
|
319
|
+
def prepare_view(view_params: Dict) -> TableClause:
|
|
320
|
+
"""Prepare SQLAlchemy view from dictionary with keys: name, metadata and
|
|
321
|
+
selectable. Must be run before Base.metadata.create_all.
|
|
322
|
+
|
|
323
|
+
Args:
|
|
324
|
+
view_params (Dict): Dictionary with keys name, metadata, selectable
|
|
325
|
+
|
|
326
|
+
Returns:
|
|
327
|
+
TableClause: SQLAlchemy View
|
|
328
|
+
"""
|
|
329
|
+
return view(**view_params)
|
|
330
|
+
|
|
331
|
+
@classmethod
|
|
332
|
+
def prepare_views(cls, view_params_list: List[Dict]) -> List[TableClause]:
|
|
333
|
+
"""Prepare SQLAlchemy views from a list of dictionaries with keys: name,
|
|
334
|
+
metadata and selectable. Must be run before Base.metadata.create_all.
|
|
335
|
+
|
|
336
|
+
Args:
|
|
337
|
+
view_params_list (List[Dict]): List of dictionaries with view parameters
|
|
338
|
+
|
|
339
|
+
Returns:
|
|
340
|
+
List[TableClause]: SQLAlchemy Views
|
|
341
|
+
"""
|
|
342
|
+
results = []
|
|
343
|
+
for view_params in view_params_list:
|
|
344
|
+
results.append(cls.prepare_view(view_params))
|
|
345
|
+
return results
|
hdx/database/views.py
CHANGED
|
@@ -7,10 +7,8 @@ https://github.com/sqlalchemy/sqlalchemy/wiki/Views#sqlalchemy-14-20-version
|
|
|
7
7
|
"""
|
|
8
8
|
|
|
9
9
|
import sqlalchemy as sa
|
|
10
|
-
from sqlalchemy import MetaData, Selectable, TableClause
|
|
11
10
|
from sqlalchemy.ext import compiler
|
|
12
11
|
from sqlalchemy.schema import DDLElement
|
|
13
|
-
from sqlalchemy.sql import table
|
|
14
12
|
|
|
15
13
|
|
|
16
14
|
class CreateView(DDLElement):
|
|
@@ -45,22 +43,15 @@ def view_doesnt_exist(ddl, target, connection, **kw):
|
|
|
45
43
|
return not view_exists(ddl, target, connection, **kw)
|
|
46
44
|
|
|
47
45
|
|
|
48
|
-
def view(name
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
Returns:
|
|
57
|
-
TableClause: SQLAlchemy View
|
|
58
|
-
"""
|
|
59
|
-
t = table(name)
|
|
60
|
-
|
|
61
|
-
t._columns._populate_separate_keys(
|
|
62
|
-
col._make_proxy(t) for col in selectable.selected_columns
|
|
46
|
+
def view(name, metadata, selectable):
|
|
47
|
+
t = sa.table(
|
|
48
|
+
name,
|
|
49
|
+
*(
|
|
50
|
+
sa.Column(c.name, c.type, primary_key=c.primary_key)
|
|
51
|
+
for c in selectable.selected_columns
|
|
52
|
+
),
|
|
63
53
|
)
|
|
54
|
+
t.primary_key.update(c for c in t.c if c.primary_key)
|
|
64
55
|
|
|
65
56
|
sa.event.listen(
|
|
66
57
|
metadata,
|
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: hdx-python-database
|
|
3
|
-
Version: 1.3.
|
|
3
|
+
Version: 1.3.6
|
|
4
4
|
Summary: HDX Python database utilities
|
|
5
5
|
Project-URL: Homepage, https://github.com/OCHA-DAP/hdx-python-database
|
|
6
6
|
Author-email: Michael Rans <rans@email.com>
|
|
@@ -28,6 +28,8 @@ Requires-Python: >=3.8
|
|
|
28
28
|
Requires-Dist: sqlalchemy>=2
|
|
29
29
|
Provides-Extra: dev
|
|
30
30
|
Requires-Dist: pre-commit; extra == 'dev'
|
|
31
|
+
Provides-Extra: docs
|
|
32
|
+
Requires-Dist: mkapi; extra == 'docs'
|
|
31
33
|
Provides-Extra: postgresql
|
|
32
34
|
Requires-Dist: psycopg[binary]; extra == 'postgresql'
|
|
33
35
|
Provides-Extra: sshtunnel
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
hdx/database/__init__.py,sha256=sG2OYV5LHskgiDkB4IeVeQyPHL22Ta1kuxf-xN92-5E,155
|
|
2
|
+
hdx/database/_version.py,sha256=O6KrLlaUs-R3ppT6iunhqb7DyTGig2RAYgNqUiGEk0w,511
|
|
3
|
+
hdx/database/database.py,sha256=hXjdRraeFouQ1afy9vHpCUkyaP-I2tPFGkDXE1z8A0A,12403
|
|
4
|
+
hdx/database/dburi.py,sha256=rw8RuLpU85vxYl5-Zk1SdZRAz-tqre_ot5aEnkjhxAQ,2977
|
|
5
|
+
hdx/database/no_timezone.py,sha256=Lw087G3qOf8H70o8Hq_VvOMZ3ANei-lIA_sUlpFU_Uo,1130
|
|
6
|
+
hdx/database/postgresql.py,sha256=i5s_CJQZagMTuzgmnLakCOnSnvczpgJ26Q1CV6lz2_A,2646
|
|
7
|
+
hdx/database/utils.py,sha256=6lLymWriMFCq8BZgrDt4V5z14GR3PbvG_WXZeBc99Ws,415
|
|
8
|
+
hdx/database/views.py,sha256=U4BN9tbjlr93qKkeNbZNo_PCIwGsng09mQ2xFqh5CH0,1636
|
|
9
|
+
hdx/database/with_timezone.py,sha256=GUUylmTwQHp22yedmCVKpVJzvjWgdS9lg9KeqciNJkY,235
|
|
10
|
+
hdx_python_database-1.3.6.dist-info/METADATA,sha256=cPzaaLmmPam1abq5Ebh-Jed_qgv3jw-4FOnzLTc77f8,2750
|
|
11
|
+
hdx_python_database-1.3.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
12
|
+
hdx_python_database-1.3.6.dist-info/licenses/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
|
|
13
|
+
hdx_python_database-1.3.6.dist-info/RECORD,,
|
|
@@ -1,12 +0,0 @@
|
|
|
1
|
-
hdx/database/__init__.py,sha256=KS5VOdLv9P-YDTE_2ZUtVm6QNwezNpXMSd7SvMYU4EM,11655
|
|
2
|
-
hdx/database/_version.py,sha256=q-rMkBhc2cbe8SdkcXn0kh8Fvtl2cIlapH0EvPC1bmc,411
|
|
3
|
-
hdx/database/dburi.py,sha256=rw8RuLpU85vxYl5-Zk1SdZRAz-tqre_ot5aEnkjhxAQ,2977
|
|
4
|
-
hdx/database/no_timezone.py,sha256=Lw087G3qOf8H70o8Hq_VvOMZ3ANei-lIA_sUlpFU_Uo,1130
|
|
5
|
-
hdx/database/postgresql.py,sha256=i5s_CJQZagMTuzgmnLakCOnSnvczpgJ26Q1CV6lz2_A,2646
|
|
6
|
-
hdx/database/utils.py,sha256=6lLymWriMFCq8BZgrDt4V5z14GR3PbvG_WXZeBc99Ws,415
|
|
7
|
-
hdx/database/views.py,sha256=R_1Ni0U2CL3IlydxcB1vyZp7QoDZim6g-807jQIaWQU,1950
|
|
8
|
-
hdx/database/with_timezone.py,sha256=GUUylmTwQHp22yedmCVKpVJzvjWgdS9lg9KeqciNJkY,235
|
|
9
|
-
hdx_python_database-1.3.4.dist-info/METADATA,sha256=HjOmHAG50Rstq_1SbBcfg7rtjoPG64L2VHFRWZwYsMU,2691
|
|
10
|
-
hdx_python_database-1.3.4.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
|
|
11
|
-
hdx_python_database-1.3.4.dist-info/licenses/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
|
|
12
|
-
hdx_python_database-1.3.4.dist-info/RECORD,,
|
{hdx_python_database-1.3.4.dist-info → hdx_python_database-1.3.6.dist-info}/licenses/LICENSE
RENAMED
|
File without changes
|