hdx-python-database 1.2.9__tar.gz → 1.3.1__tar.gz

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.
Files changed (36) hide show
  1. hdx_python_database-1.3.1/.config/pre-commit-config.yaml +26 -0
  2. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/.config/ruff.toml +7 -5
  3. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/CONTRIBUTING.md +2 -3
  4. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/PKG-INFO +3 -4
  5. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/README.md +1 -2
  6. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/documentation/.readthedocs.yaml +1 -0
  7. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/documentation/main.md +35 -8
  8. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/documentation/pydoc-markdown.yaml +1 -1
  9. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/pyproject.toml +4 -12
  10. hdx_python_database-1.3.1/requirements.txt +57 -0
  11. hdx_python_database-1.3.1/src/hdx/database/__init__.py +294 -0
  12. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/src/hdx/database/_version.py +2 -2
  13. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/src/hdx/database/dburi.py +2 -0
  14. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/src/hdx/database/no_timezone.py +1 -0
  15. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/src/hdx/database/postgresql.py +1 -0
  16. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/src/hdx/database/views.py +3 -31
  17. hdx_python_database-1.3.1/tests/hdx/database/conftest.py +86 -0
  18. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/tests/hdx/database/dbtestdate.py +2 -5
  19. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/tests/hdx/database/test_database.py +51 -44
  20. hdx_python_database-1.2.9/.config/black.toml +0 -2
  21. hdx_python_database-1.2.9/.config/pre-commit-config.yaml +0 -26
  22. hdx_python_database-1.2.9/requirements.txt +0 -69
  23. hdx_python_database-1.2.9/src/hdx/database/__init__.py +0 -149
  24. hdx_python_database-1.2.9/tests/hdx/database/conftest.py +0 -22
  25. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/.config/coveragerc +0 -0
  26. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/.config/pytest.ini +0 -0
  27. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/.github/workflows/publish.yaml +0 -0
  28. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/.github/workflows/run-python-tests.yaml +0 -0
  29. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/.gitignore +0 -0
  30. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/LICENSE +0 -0
  31. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/src/hdx/database/utils.py +0 -0
  32. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/src/hdx/database/with_timezone.py +0 -0
  33. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/tests/fixtures/test.db +0 -0
  34. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/tests/hdx/database/__init__.py +0 -0
  35. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/tests/hdx/database/test_dburi.py +0 -0
  36. {hdx_python_database-1.2.9 → hdx_python_database-1.3.1}/tests/hdx/database/test_postgresql.py +0 -0
@@ -0,0 +1,26 @@
1
+ default_language_version:
2
+ python: python3.11
3
+ repos:
4
+ - repo: https://github.com/pre-commit/pre-commit-hooks
5
+ rev: v4.5.0
6
+ hooks:
7
+ - id: trailing-whitespace
8
+ - id: end-of-file-fixer
9
+ - id: check-ast
10
+ - repo: https://github.com/astral-sh/ruff-pre-commit
11
+ rev: v0.4.2
12
+ hooks:
13
+ # Run the linter.
14
+ - id: ruff
15
+ args: [--config, .config/ruff.toml, --fix]
16
+ # Run the formatter.
17
+ - id: ruff-format
18
+ args: [--config, .config/ruff.toml]
19
+ - repo: https://github.com/astral-sh/uv-pre-commit
20
+ rev: 0.1.38
21
+ hooks:
22
+ # Run the pip compile
23
+ - id: pip-compile
24
+ name: pip-compile requirements.txt
25
+ files: pyproject.toml
26
+ args: [ pyproject.toml, --resolver=backtracking, --all-extras, --upgrade, -q, -o, requirements.txt ]
@@ -1,14 +1,16 @@
1
1
  line-length = 79
2
2
  exclude = ["_version.py"]
3
- ignore = [
4
- "E501" # Line too long
5
- ]
6
- # List of rules: https://beta.ruff.rs/docs/rules/
3
+
4
+ [lint]
5
+ # List of rules: https://docs.astral.sh/ruff/rules/
7
6
  select = [
8
7
  "E", # pycodestyle - default
9
8
  "F", # pyflakes - default
10
9
  "I" # isort
11
10
  ]
11
+ ignore = [
12
+ "E501" # Line too long
13
+ ]
12
14
 
13
- [isort]
15
+ [lint.isort]
14
16
  known-local-folder = ["hdx"]
@@ -21,14 +21,13 @@ you make a git commit:
21
21
  pre-commit install
22
22
 
23
23
  The configuration file for this project is in a
24
- non-start location. Thus, you will need to edit your
24
+ non-standard location. Thus, you will need to edit your
25
25
  `.git/hooks/pre-commit` file to reflect this. Change
26
26
  the line that begins with `ARGS` to:
27
27
 
28
28
  ARGS=(hook-impl --config=.config/pre-commit-config.yaml --hook-type=pre-commit)
29
29
 
30
30
  With pre-commit, all code is formatted according to
31
- [black]("https://github.com/psf/black") and
32
31
  [ruff]("https://github.com/charliermarsh/ruff") guidelines.
33
32
 
34
33
  To check if your changes pass pre-commit without committing, run:
@@ -41,7 +40,7 @@ To run the tests and view coverage, execute:
41
40
 
42
41
  pytest -c .config/pytest.ini --cov hdx --cov-config .config/coveragerc
43
42
 
44
- Follow the example set out already in ``api.rst`` as you write the documentation.
43
+ Follow the example set out already in ``documentation/main.md`` as you write the documentation.
45
44
 
46
45
  ## Packages
47
46
 
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.3
2
2
  Name: hdx-python-database
3
- Version: 1.2.9
3
+ Version: 1.3.1
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>
@@ -39,8 +39,7 @@ Description-Content-Type: text/markdown
39
39
 
40
40
  [![Build Status](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yaml/badge.svg)](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yaml)
41
41
  [![Coverage Status](https://coveralls.io/repos/github/OCHA-DAP/hdx-python-database/badge.svg?branch=main&ts=1)](https://coveralls.io/github/OCHA-DAP/hdx-python-database?branch=main)
42
- [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
43
- [![Imports: isort](https://img.shields.io/badge/%20imports-isort-%231674b1?style=flat&labelColor=ef8336)](https://pycqa.github.io/isort/)
42
+ [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
44
43
  [![Downloads](https://img.shields.io/pypi/dm/hdx-python-database.svg)](https://pypistats.org/packages/hdx-python-database)
45
44
 
46
45
  The HDX Python Database Library provides utilities for connecting to databases in a standard way including
@@ -1,7 +1,6 @@
1
1
  [![Build Status](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yaml/badge.svg)](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yaml)
2
2
  [![Coverage Status](https://coveralls.io/repos/github/OCHA-DAP/hdx-python-database/badge.svg?branch=main&ts=1)](https://coveralls.io/github/OCHA-DAP/hdx-python-database?branch=main)
3
- [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
4
- [![Imports: isort](https://img.shields.io/badge/%20imports-isort-%231674b1?style=flat&labelColor=ef8336)](https://pycqa.github.io/isort/)
3
+ [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
5
4
  [![Downloads](https://img.shields.io/pypi/dm/hdx-python-database.svg)](https://pypistats.org/packages/hdx-python-database)
6
5
 
7
6
  The HDX Python Database Library provides utilities for connecting to databases in a standard way including
@@ -11,5 +11,6 @@ build:
11
11
  jobs:
12
12
  pre_build:
13
13
  - pip install --upgrade mkdocs
14
+ - pip install mkdocs-material
14
15
  - pip install pydoc-markdown
15
16
  - pydoc-markdown documentation/pydoc-markdown.yaml
@@ -22,6 +22,13 @@ installed with extra `postgresql`:
22
22
  pip install hdx-python-database[postgresql]
23
23
 
24
24
  ## Breaking changes
25
+ From 1.3.1, Database class refactored. With returns Database not Session
26
+ object and can accept a prepare function called before
27
+ Base.metadata.create_all. get_session can be used to obtain the session.
28
+ get_engine, get_reflected_classes, get_prepare_results methods added.
29
+ build_view, build_views move into Database class and renamed prepare_view and
30
+ prepare_views.
31
+
25
32
  From 1.2.8, the sshtunnel dependency is optional.
26
33
 
27
34
  From 1.2.7, default table names are no longer plural. The camel case class name
@@ -51,8 +58,13 @@ Versions from 1.0.6 no longer support Python 2.7.
51
58
 
52
59
  ## Database
53
60
 
61
+ A fresh schema can be created using the function `recreate_schema`. Any
62
+ existing schema with that name will be dropped eg.
63
+
64
+ Database.recreate_schema(engine, db_uri)
65
+
54
66
  Your SQLAlchemy database tables must inherit from `Base` in
55
- `hdx.database.no_timezone` or `hdx.database.with_timezone` eg. :
67
+ `hdx.database.no_timezone` or `hdx.database.with_timezone` eg.
56
68
 
57
69
  from hdx.database.no_timezone import Base
58
70
  class MyTable(Base):
@@ -70,9 +82,14 @@ tunnel (which requires installing `hdx-python-database[sshtunnel]`):
70
82
  from hdx.database import Database
71
83
  with Database(database="db", host="1.2.3.4", username="user",
72
84
  password="pass", dialect="dialect", driver="driver",
85
+
73
86
  ssh_host="5.6.7.8", ssh_port=2222, ssh_username="sshuser",
74
87
  ssh_private_key="path_to_key", db_has_tz=True,
75
- reflect=False) as session:
88
+ reflect=True, prepare_fn=prepare_views) as database:
89
+ engine = database.get_engine()
90
+ session = database.get_session()
91
+ reflected_classes = database.get_reflected_classes()
92
+ views = database.get_prepare_results()
76
93
  session.query(...)
77
94
 
78
95
  `db_has_tz` which defaults to `False` indicates whether database datetime
@@ -83,11 +100,20 @@ between Python datetimes with timezones to timezoneless database columns.
83
100
 
84
101
  If `reflect` (which defaults to `False`) is `True`, classes will be reflected
85
102
  from an existing database and the reflected classes are returned in a variable
86
- `reflected_classes` in the returned Session object. Note that type annotation
103
+ `reflected_classes` in the Database object. Note that type annotation
87
104
  maps don't work with reflection and hence `db_has_tz` will be ignored ie.
88
105
  there will be no conversion between Python datetimes with timezones to
89
106
  timezoneless database columns.
90
107
 
108
+ There is an option to wipe and create an empty schema in the database by
109
+ setting `recreate_schema` to `True` and setting a `schema_name` ("public" is
110
+ the default).
111
+
112
+ If a prepare function is supplied in prepare_fn, it will be executed before
113
+ Base.metadata.create_all and the results of it returned in instance variable
114
+ prepare_results.
115
+
116
+
91
117
  ## Connection URI
92
118
 
93
119
  There are functions to handle converting from connection URIs to parameters and
@@ -117,10 +143,9 @@ URI that contains both dialect and driver.
117
143
 
118
144
  ## Views
119
145
 
120
- The method to make views described [here](https://github.com/sqlalchemy/sqlalchemy/wiki/Views#sqlalchemy-14-20-version)
121
- is available in this library.
146
+ The method to prepare views described [here](https://github.com/sqlalchemy/sqlalchemy/wiki/Views#sqlalchemy-14-20-version) is available in this library.
122
147
 
123
- This allows creating views like this:
148
+ This allows preparing views like this:
124
149
  ```
125
150
  class DBOrgType(Base):
126
151
  __tablename__ = "org_type"
@@ -144,11 +169,13 @@ date_view_params = {
144
169
  "selectable": select(*DBTestDate.__table__.columns),
145
170
  }
146
171
 
147
- date_view = build_view(date_view_params)
172
+ date_view = Database.prepare_view(date_view_params)
148
173
 
149
- date_view = build_views([date_view_params])[0]
174
+ date_view = Database.prepare_views([date_view_params])[0]
150
175
  ```
151
176
 
177
+ View preparation must be done before calling Base.metadata.create_all.
178
+
152
179
  ## PostgreSQL specific
153
180
 
154
181
  There is a PostgreSQL specific call that only returns when the PostgreSQL database
@@ -9,7 +9,7 @@ renderer:
9
9
  output_directory: docs
10
10
  mkdocs_config:
11
11
  site_name: HDX Python Database
12
- theme: mkdocs
12
+ theme: material
13
13
  repo_url: "https://github.com/OCHA-DAP/hdx-python-database"
14
14
  markdown:
15
15
  source_linker:
@@ -86,22 +86,14 @@ test = """
86
86
  """
87
87
 
88
88
  [[tool.hatch.envs.test.matrix]]
89
- python = ["3.11"]
89
+ python = ["3.12"]
90
90
 
91
91
  [tool.hatch.envs.lint]
92
92
  detached = true
93
- dependencies = [
94
- "black",
95
- "ruff",
96
- ]
93
+ dependencies = ["ruff"]
97
94
 
98
95
  [tool.hatch.envs.lint.scripts]
99
96
  style = [
100
- "ruff --config .config/ruff.toml {args:.}",
101
- "black --config .config/black.toml --check --diff {args:.}",
102
- ]
103
- # Not used for anything at the moment
104
- fmt = [
105
- "black --config .config/black.toml {args:.}",
106
- "ruff --config .config/ruff.toml --fix {args:.}",
97
+ "ruff check --config .config/ruff.toml --diff {args:.}",
98
+ "ruff format --config .config/ruff.toml --diff {args:.}",
107
99
  ]
@@ -0,0 +1,57 @@
1
+ # This file was autogenerated by uv via the following command:
2
+ # uv pip compile pyproject.toml --resolver=backtracking --all-extras -o requirements.txt
3
+ bcrypt==4.1.3
4
+ # via paramiko
5
+ cffi==1.16.0
6
+ # via
7
+ # cryptography
8
+ # pynacl
9
+ cfgv==3.4.0
10
+ # via pre-commit
11
+ coverage==7.5.1
12
+ # via pytest-cov
13
+ cryptography==42.0.7
14
+ # via paramiko
15
+ distlib==0.3.8
16
+ # via virtualenv
17
+ filelock==3.14.0
18
+ # via virtualenv
19
+ greenlet==3.0.3
20
+ # via sqlalchemy
21
+ identify==2.5.36
22
+ # via pre-commit
23
+ iniconfig==2.0.0
24
+ # via pytest
25
+ nodeenv==1.8.0
26
+ # via pre-commit
27
+ packaging==24.0
28
+ # via pytest
29
+ paramiko==3.4.0
30
+ # via sshtunnel
31
+ platformdirs==4.2.1
32
+ # via virtualenv
33
+ pluggy==1.5.0
34
+ # via pytest
35
+ pre-commit==3.7.0
36
+ psycopg==3.1.18
37
+ psycopg-binary==3.1.18
38
+ # via psycopg
39
+ pycparser==2.22
40
+ # via cffi
41
+ pynacl==1.5.0
42
+ # via paramiko
43
+ pytest==8.2.0
44
+ # via pytest-cov
45
+ pytest-cov==5.0.0
46
+ pyyaml==6.0.1
47
+ # via pre-commit
48
+ setuptools==69.5.1
49
+ # via nodeenv
50
+ sqlalchemy==2.0.30
51
+ sshtunnel==0.4.0
52
+ typing-extensions==4.11.0
53
+ # via
54
+ # psycopg
55
+ # sqlalchemy
56
+ virtualenv==20.26.1
57
+ # via pre-commit
@@ -0,0 +1,294 @@
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, sessionmaker
10
+ from sqlalchemy.pool import NullPool
11
+ from sqlalchemy.sql.ddl import CreateSchema, DropSchema
12
+ from sqlalchemy.util import Properties
13
+
14
+ from ._version import version as __version__ # noqa: F401
15
+ from .dburi import get_connection_uri
16
+ from .no_timezone import Base as NoTZBase
17
+ from .postgresql import wait_for_postgresql
18
+ from .views import view
19
+ from .with_timezone import Base as TZBase
20
+
21
+ logger = logging.getLogger(__name__)
22
+
23
+
24
+ class DatabaseError(Exception):
25
+ pass
26
+
27
+
28
+ def do_nothing_fn():
29
+ return None
30
+
31
+
32
+ class Database:
33
+ """Database helper class to handle ssh tunnels, waiting for PostgreSQL to
34
+ be up etc. Can be used in a with statement returning a Database object that
35
+ if reflect is True will have a variable reflected_classes containing the
36
+ reflected classes. Either an SQLAlchemy engine, database uri or various
37
+ database parameters must be supplied.
38
+
39
+ db_has_tz which defaults to False indicates whether database datetime
40
+ columns have timezones. If not, conversion occurs between Python datetimes
41
+ with timezones to timezoneless database columns (but not when using
42
+ reflection). If table_base is supplied, db_has_tz is ignored.
43
+
44
+ There is an option to wipe and create an empty schema in the database by
45
+ setting recreate_schema to True and setting a schema_name ("public" is the
46
+ default).
47
+
48
+ If a prepare function is supplied in prepare_fn, it will be executed before
49
+ Base.metadata.create_all and the results of it returned in instance variable
50
+ prepare_results.
51
+
52
+ Args:
53
+ engine (Optional[Engine]): SQLAlchemy engine to use.
54
+ db_uri (Optional[str]): Connection URI.
55
+ database (Optional[str]): Database name
56
+ host (Optional[str]): Host where database is located
57
+ port (Union[int, str, None]): Database port
58
+ username (Optional[str]): Username to log into database
59
+ password (Optional[str]): Password to log into database
60
+ dialect (str): Database dialect. Defaults to "postgresql".
61
+ driver (Optional[str]): Database driver. Defaults to None (psycopg if postgresql or None)
62
+ db_has_tz (bool): True if db datetime columns have timezone. Defaults to False.
63
+ table_base (Optional[Type[DeclarativeBase]]): Override table base. Defaults to None.
64
+ reflect (bool): Whether to reflect existing tables. Defaults to False.
65
+ **kwargs: See below
66
+ recreate_schema (bool): Whether to recreate schema
67
+ schema_name (str): Database schema name. Defaults to "public".
68
+ prepare_fn (Callable[[], None]]): Function to call before Base.metadata.create_all.
69
+ ssh_host (str): SSH host (the server to connect to)
70
+ ssh_port (int): SSH port. Defaults to 22.
71
+ ssh_username (str): SSH username
72
+ ssh_password (str): SSH password
73
+ ssh_private_key: Location of SSH private key (instead of ssh_password)
74
+ For more advanced usage, see SSHTunnelForwarder documentation.
75
+
76
+ """
77
+
78
+ def __init__(
79
+ self,
80
+ engine: Optional[Engine] = None,
81
+ db_uri: Optional[str] = None,
82
+ database: Optional[str] = None,
83
+ host: Optional[str] = None,
84
+ port: Union[int, str, None] = None,
85
+ username: Optional[str] = None,
86
+ password: Optional[str] = None,
87
+ dialect: str = "postgresql",
88
+ driver: Optional[str] = None,
89
+ db_has_tz: bool = False,
90
+ table_base: Optional[Type[DeclarativeBase]] = None,
91
+ reflect: bool = False,
92
+ **kwargs: Any,
93
+ ) -> None:
94
+ if port is not None:
95
+ port = int(port)
96
+ schema_name = None
97
+ if len(kwargs) == 0:
98
+ recreate_schema = False
99
+ prepare_fn = do_nothing_fn
100
+ else:
101
+ recreate_schema = kwargs.pop("recreate_schema", False)
102
+ schema_name = kwargs.pop("schema", "public")
103
+ prepare_fn = kwargs.pop("prepare_fn", do_nothing_fn)
104
+ if len(kwargs) != 0:
105
+ try:
106
+ import sshtunnel
107
+ except ImportError:
108
+ # dependency missing, log an error
109
+ logger.error(
110
+ "sshtunnel not found! Please install hdx-python-database[sshtunnel] to enable."
111
+ )
112
+ raise
113
+ ssh_host = kwargs["ssh_host"]
114
+ del kwargs["ssh_host"]
115
+ ssh_port = kwargs.get("ssh_port")
116
+ if ssh_port is not None:
117
+ ssh_port = int(ssh_port)
118
+ del kwargs["ssh_port"]
119
+ else:
120
+ ssh_port = 22
121
+ self.server = sshtunnel.SSHTunnelForwarder(
122
+ (ssh_host, ssh_port),
123
+ remote_bind_address=(host, port),
124
+ **kwargs,
125
+ )
126
+ self.server.start()
127
+ host = self.server.local_bind_host
128
+ port = self.server.local_bind_port
129
+ else:
130
+ self.server = None
131
+ if not engine and not db_uri:
132
+ db_uri = get_connection_uri(
133
+ database,
134
+ host,
135
+ port,
136
+ username,
137
+ password,
138
+ dialect=dialect,
139
+ driver=driver,
140
+ )
141
+ if not db_uri:
142
+ raise DatabaseError(
143
+ "No engine, database uri or connection parameters provided!"
144
+ )
145
+ if not engine and dialect == "postgresql":
146
+ wait_for_postgresql(db_uri)
147
+ if not table_base:
148
+ if db_has_tz:
149
+ table_base = TZBase
150
+ else:
151
+ table_base = NoTZBase
152
+ if not engine:
153
+ engine = create_engine(db_uri, poolclass=NullPool, echo=False)
154
+ self.engine: Engine = engine
155
+ if recreate_schema:
156
+ self.recreate_schema(engine, schema_name)
157
+ self.prepare_results = prepare_fn()
158
+ self.session, self.reflected_classes = self.create_session(
159
+ engine,
160
+ table_base=table_base,
161
+ reflect=reflect,
162
+ )
163
+
164
+ def __enter__(self) -> "Database":
165
+ return self
166
+
167
+ def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
168
+ self.session.close()
169
+ if self.server is not None:
170
+ self.server.stop()
171
+
172
+ def get_engine(self) -> Engine:
173
+ """Returns SQLAlchemy engine.
174
+
175
+ Returns:
176
+ sqlalchemy.Engine: SQLAlchemy engine
177
+ """
178
+ return self.engine
179
+
180
+ def get_session(self) -> Session:
181
+ """Returns SQLAlchemy session.
182
+
183
+ Returns:
184
+ sqlalchemy.orm.Session: SQLAlchemy session
185
+ """
186
+ return self.session
187
+
188
+ def get_reflected_classes(self) -> Any:
189
+ """Gets reflected classes.
190
+
191
+ Returns:
192
+ Any: Reflected classes
193
+ """
194
+ return self.reflected_classes
195
+
196
+ def get_prepare_results(self) -> Any:
197
+ """Returns results from prepare function.
198
+
199
+ Returns:
200
+ Any: Results from prepare function
201
+ """
202
+ return self.prepare_results
203
+
204
+ @staticmethod
205
+ def create_session(
206
+ engine: Optional[Engine] = None,
207
+ db_uri: Optional[str] = None,
208
+ table_base: Type[DeclarativeBase] = NoTZBase,
209
+ reflect: bool = False,
210
+ ) -> Tuple[Session, Optional[Properties]]:
211
+ """Creates SQLAlchemy session given SQLAlchemy engine or database uri
212
+ (one of which must be supplied). Tables must inherit from Base in
213
+ hdx.utilities.database unless base is defined. If reflect is True,
214
+ classes will be reflected from an existing database and the reflected
215
+ classes will be returned. Note that type annotation maps don't work
216
+ with reflection.
217
+
218
+ Args:
219
+ engine (Optional[Engine]): SQLAlchemy engine to use. Defaults to None (create from db_uri).
220
+ db_uri (Optional[str]): Connection URI. Defaults to None (use engine).
221
+ table_base (Type[DeclarativeBase]): Base database table class. Defaults to NoTZBase.
222
+ reflect (bool): Whether to reflect existing tables. Defaults to False.
223
+
224
+ Returns:
225
+ Tuple[Session, Optional[Properties]]: (SQLAlchemy session, reflected classes if available)
226
+ """
227
+ if not engine:
228
+ if db_uri is None:
229
+ raise DatabaseError("No engine or database uri provided!")
230
+ engine = create_engine(db_uri, poolclass=NullPool, echo=False)
231
+ Session = sessionmaker(bind=engine)
232
+ if reflect:
233
+ Base = automap_base(declarative_base=table_base)
234
+ Base.prepare(autoload_with=engine)
235
+ reflected_classes = Base.classes
236
+ else:
237
+ table_base.metadata.create_all(engine)
238
+ reflected_classes = None
239
+ session = Session()
240
+ return session, reflected_classes
241
+
242
+ @staticmethod
243
+ def recreate_schema(engine: Engine, schema_name: str = "public") -> bool:
244
+ """Wipe and create empty schema in database using SQLAlchemy.
245
+
246
+ Args:
247
+ engine (Engine): SQLAlchemy engine to use.
248
+ schema_name (str): Schema name. Defaults to "public".
249
+
250
+ Returns:
251
+ bool: True if all successful, False if not
252
+ """
253
+ # Wipe and create an empty schema
254
+ try:
255
+ with engine.connect() as connection:
256
+ connection.execute(
257
+ DropSchema(schema_name, cascade=True, if_exists=True)
258
+ )
259
+ connection.commit()
260
+ connection.execute(
261
+ CreateSchema(schema_name, if_not_exists=True)
262
+ )
263
+ connection.commit()
264
+ return True
265
+ except SQLAlchemyError:
266
+ return False
267
+
268
+ @staticmethod
269
+ def prepare_view(view_params: Dict) -> TableClause:
270
+ """Prepare SQLAlchemy view from dictionary with keys: name, metadata and
271
+ selectable. Must be run before Base.metadata.create_all.
272
+
273
+ Args:
274
+ view_params (Dict): Dictionary with keys name, metadata, selectable
275
+
276
+ Returns:
277
+ TableClause: SQLAlchemy View
278
+ """
279
+ return view(**view_params)
280
+
281
+ @classmethod
282
+ def prepare_views(cls, view_params_list: List[Dict]) -> List[TableClause]:
283
+ """Prepare SQLAlchemy views from a list of dictionaries with keys: name,
284
+ metadata and selectable. Must be run before Base.metadata.create_all.
285
+
286
+ Args:
287
+ view_params_list (List[Dict]): List of dictionaries with view parameters
288
+ Returns:
289
+ List[TableClause]: SQLAlchemy Views
290
+ """
291
+ results = []
292
+ for view_params in view_params_list:
293
+ results.append(cls.prepare_view(view_params))
294
+ return results
@@ -12,5 +12,5 @@ __version__: str
12
12
  __version_tuple__: VERSION_TUPLE
13
13
  version_tuple: VERSION_TUPLE
14
14
 
15
- __version__ = version = '1.2.9'
16
- __version_tuple__ = version_tuple = (1, 2, 9)
15
+ __version__ = version = '1.3.1'
16
+ __version_tuple__ = version_tuple = (1, 3, 1)
@@ -62,6 +62,8 @@ def get_connection_uri(
62
62
  Returns:
63
63
  str: Connection URI
64
64
  """
65
+ if not database:
66
+ return ""
65
67
  strings = [dialect]
66
68
  if include_driver:
67
69
  if dialect == "postgresql" and driver is None:
@@ -1,4 +1,5 @@
1
1
  """Utilities for database datetime columns without timezone"""
2
+
2
3
  from datetime import datetime, timezone
3
4
  from typing import Optional
4
5
 
@@ -1,4 +1,5 @@
1
1
  """PostgreSQL specific utilities"""
2
+
2
3
  import logging
3
4
  import time
4
5
 
@@ -1,10 +1,10 @@
1
1
  """
2
- Create a DB View.
2
+ Prepare database views ins SQLAlchemy. Must be run before calling
3
+ Base.metadata.create_all.
3
4
 
4
5
  Copied from:
5
6
  https://github.com/sqlalchemy/sqlalchemy/wiki/Views#sqlalchemy-14-20-version
6
7
  """
7
- from typing import Dict, List
8
8
 
9
9
  import sqlalchemy as sa
10
10
  from sqlalchemy import MetaData, Selectable, TableClause
@@ -46,7 +46,7 @@ def view_doesnt_exist(ddl, target, connection, **kw):
46
46
 
47
47
 
48
48
  def view(name: str, metadata: MetaData, selectable: Selectable) -> TableClause:
49
- """Create SQLAlchemy view
49
+ """Prepare SQLAlchemy view. Must be run before Base.metadata.create_all.
50
50
 
51
51
  Args:
52
52
  name (str): View name
@@ -73,31 +73,3 @@ def view(name: str, metadata: MetaData, selectable: Selectable) -> TableClause:
73
73
  DropView(name).execute_if(callable_=view_exists),
74
74
  )
75
75
  return t
76
-
77
-
78
- def build_view(view_params: Dict) -> TableClause:
79
- """Create SQLAlchemy view from dictionary with keys: name, metadata and
80
- selectable
81
-
82
- Args:
83
- view_params (Dict): Dictionary with keys name, metadata, selectable
84
-
85
- Returns:
86
- TableClause: SQLAlchemy View
87
- """
88
- return view(**view_params)
89
-
90
-
91
- def build_views(view_params_list: List[Dict]) -> List[TableClause]:
92
- """Create SQLAlchemy views from a list of dictionaries with keys: name,
93
- metadata and selectable
94
-
95
- Args:
96
- view_params_list (List[Dict]): List of dictionaries with view parameters
97
- Returns:
98
- List[TableClause]: SQLAlchemy Views
99
- """
100
- results = []
101
- for view_params in view_params_list:
102
- results.append(build_view(view_params))
103
- return results
@@ -0,0 +1,86 @@
1
+ from collections import namedtuple
2
+
3
+ import psycopg
4
+ import pytest
5
+ from sqlalchemy.exc import SQLAlchemyError
6
+ from sshtunnel import SSHTunnelForwarder
7
+
8
+ from . import PsycopgConnection
9
+ from .test_database import TestDatabase
10
+ from hdx.database import Database
11
+
12
+
13
+ @pytest.fixture(scope="function")
14
+ def mock_psycopg(monkeypatch):
15
+ def connect(*args, **kwargs):
16
+ if PsycopgConnection.connected:
17
+
18
+ class PGConnection:
19
+ @staticmethod
20
+ def close():
21
+ return
22
+
23
+ return PGConnection()
24
+ else:
25
+ PsycopgConnection.connected = True
26
+ raise psycopg.OperationalError
27
+
28
+ monkeypatch.setattr(psycopg, "connect", connect)
29
+
30
+
31
+ @pytest.fixture(scope="function")
32
+ def mock_SSHTunnelForwarder(monkeypatch):
33
+ def init(*args, **kwargs):
34
+ return None
35
+
36
+ def start(_):
37
+ TestDatabase.started = True
38
+
39
+ def stop(_):
40
+ TestDatabase.stopped = True
41
+
42
+ monkeypatch.setattr(SSHTunnelForwarder, "__init__", init)
43
+ monkeypatch.setattr(SSHTunnelForwarder, "start", start)
44
+ monkeypatch.setattr(SSHTunnelForwarder, "stop", stop)
45
+ monkeypatch.setattr(SSHTunnelForwarder, "local_bind_host", "0.0.0.0")
46
+ monkeypatch.setattr(SSHTunnelForwarder, "local_bind_port", 12345)
47
+
48
+ def create_session(_, engine, table_base, reflect):
49
+ class Session:
50
+ bind = namedtuple("Bind", "engine")
51
+
52
+ def close(self):
53
+ return None
54
+
55
+ Session.bind.engine = engine
56
+ TestDatabase.table_base = table_base
57
+ return Session(), None
58
+
59
+ monkeypatch.setattr(Database, "create_session", create_session)
60
+
61
+
62
+ @pytest.fixture(scope="function")
63
+ def mock_engine():
64
+ class SQAConnection:
65
+ def __enter__(self):
66
+ return self
67
+
68
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
69
+ pass
70
+
71
+ @staticmethod
72
+ def execute(statement):
73
+ if "mydatabase" in str(statement):
74
+ return
75
+ raise SQLAlchemyError("Fail!")
76
+
77
+ @staticmethod
78
+ def commit():
79
+ return
80
+
81
+ class MockEngine:
82
+ @staticmethod
83
+ def connect():
84
+ return SQAConnection()
85
+
86
+ return MockEngine()
@@ -1,12 +1,11 @@
1
- """SQLAlchemy class representing DBTestDate row. Holds test data for dates.
2
- """
1
+ """SQLAlchemy class representing DBTestDate row. Holds test data for dates."""
2
+
3
3
  from datetime import datetime
4
4
 
5
5
  from sqlalchemy import select
6
6
  from sqlalchemy.orm import Mapped, mapped_column
7
7
 
8
8
  from hdx.database.no_timezone import Base
9
- from hdx.database.views import build_views
10
9
 
11
10
 
12
11
  class DBTestDate(Base):
@@ -30,5 +29,3 @@ date_view_params = {
30
29
  "metadata": Base.metadata,
31
30
  "selectable": select(*DBTestDate.__table__.columns),
32
31
  }
33
-
34
- date_view = build_views([date_view_params])[0]
@@ -1,5 +1,5 @@
1
1
  """Database Utility Tests"""
2
- from collections import namedtuple
2
+
3
3
  from copy import deepcopy
4
4
  from datetime import datetime, timezone
5
5
  from os import remove
@@ -8,10 +8,9 @@ from shutil import copyfile
8
8
 
9
9
  import pytest
10
10
  from sqlalchemy import select
11
- from sshtunnel import SSHTunnelForwarder
12
11
 
13
- from .dbtestdate import DBTestDate, date_view
14
- from hdx.database import Database
12
+ from .dbtestdate import DBTestDate, date_view_params
13
+ from hdx.database import Database, DatabaseError
15
14
  from hdx.database.no_timezone import Base as NoTZBase
16
15
  from hdx.database.with_timezone import Base as TZBase
17
16
 
@@ -49,42 +48,19 @@ class TestDatabase:
49
48
  pass
50
49
  return f"sqlite:///{TestDatabase.dbpath}"
51
50
 
52
- @pytest.fixture(scope="function")
53
- def mock_SSHTunnelForwarder(self, monkeypatch):
54
- def init(*args, **kwargs):
55
- return None
56
-
57
- def start(_):
58
- TestDatabase.started = True
59
-
60
- def stop(_):
61
- TestDatabase.stopped = True
62
-
63
- monkeypatch.setattr(SSHTunnelForwarder, "__init__", init)
64
- monkeypatch.setattr(SSHTunnelForwarder, "start", start)
65
- monkeypatch.setattr(SSHTunnelForwarder, "stop", stop)
66
- monkeypatch.setattr(SSHTunnelForwarder, "local_bind_host", "0.0.0.0")
67
- monkeypatch.setattr(SSHTunnelForwarder, "local_bind_port", 12345)
68
-
69
- def get_session(_, db_url, table_base, reflect):
70
- class Session:
71
- bind = namedtuple("Bind", "engine")
72
-
73
- def close(self):
74
- return None
75
-
76
- Session.bind.engine = namedtuple("Engine", "url")
77
- Session.bind.engine.url = db_url
78
- TestDatabase.table_base = table_base
79
- return Session()
80
-
81
- monkeypatch.setattr(Database, "get_session", get_session)
82
-
83
51
  def test_get_session(self, nodatabase):
84
52
  assert DBTestDate.__tablename__ == "db_test_date"
53
+
54
+ def prepare_fn():
55
+ return Database.prepare_views([date_view_params])
56
+
85
57
  with Database(
86
- database=TestDatabase.dbpath, port=None, dialect="sqlite"
87
- ) as dbsession:
58
+ database=TestDatabase.dbpath,
59
+ port=None,
60
+ dialect="sqlite",
61
+ prepare_fn=prepare_fn,
62
+ ) as dbdatabase:
63
+ dbsession = dbdatabase.session
88
64
  assert str(dbsession.bind.engine.url) == nodatabase
89
65
  assert TestDatabase.table_base == NoTZBase
90
66
  now = datetime(2022, 10, 20, 22, 35, 55, tzinfo=timezone.utc)
@@ -94,8 +70,14 @@ class TestDatabase:
94
70
  dbsession.commit()
95
71
  dbtestdate = dbsession.execute(select(DBTestDate)).scalar_one()
96
72
  assert dbtestdate.test_date == now
73
+ date_view = dbdatabase.get_prepare_results()[0]
97
74
  dbtestdate = dbsession.execute(select(date_view)).scalar_one()
98
75
  assert dbtestdate == now
76
+ with pytest.raises(DatabaseError):
77
+ Database.create_session()
78
+ with pytest.raises(DatabaseError):
79
+ with Database():
80
+ pass
99
81
 
100
82
  def test_get_reflect_session(self, database_to_reflect):
101
83
  with Database(
@@ -103,10 +85,13 @@ class TestDatabase:
103
85
  port=None,
104
86
  dialect="sqlite",
105
87
  reflect=True,
106
- ) as dbsession:
88
+ ) as dbdatabase:
89
+ dbengine = dbdatabase.get_engine()
90
+ assert str(dbengine.url) == "sqlite:///tests/test_database.db"
91
+ dbsession = dbdatabase.get_session()
107
92
  assert TestDatabase.table_base == NoTZBase
108
93
  assert str(dbsession.bind.engine.url) == database_to_reflect
109
- Table1 = dbsession.reflected_classes.table1
94
+ Table1 = dbdatabase.get_reflected_classes().table1
110
95
  row = dbsession.execute(select(Table1)).scalar_one()
111
96
  assert row.id == "1"
112
97
  assert row.col1 == "wfrefds"
@@ -114,19 +99,25 @@ class TestDatabase:
114
99
  # we don't have a timezone here
115
100
  assert row.date1 == datetime(1993, 9, 23, 14, 12, 56, 111000)
116
101
 
117
- def test_get_session_ssh(self, mock_psycopg, mock_SSHTunnelForwarder):
102
+ def test_get_session_ssh(
103
+ self,
104
+ mock_psycopg,
105
+ mock_SSHTunnelForwarder,
106
+ ):
118
107
  with Database(
119
108
  ssh_host="mysshhost", **TestDatabase.params_pg
120
- ) as dbsession:
109
+ ) as dbdatabase:
110
+ dbsession = dbdatabase.get_session()
121
111
  assert (
122
112
  str(dbsession.bind.engine.url)
123
- == "postgresql+psycopg://myuser:mypass@0.0.0.0:12345/mydatabase"
113
+ == "postgresql+psycopg://myuser:***@0.0.0.0:12345/mydatabase"
124
114
  )
125
115
  params = deepcopy(TestDatabase.params_pg)
126
116
  del params["password"]
127
117
  with Database(
128
118
  ssh_host="mysshhost", ssh_port=25, **params
129
- ) as dbsession:
119
+ ) as dbdatabase:
120
+ dbsession = dbdatabase.get_session()
130
121
  assert (
131
122
  str(dbsession.bind.engine.url)
132
123
  == "postgresql+psycopg://myuser@0.0.0.0:12345/mydatabase"
@@ -134,9 +125,25 @@ class TestDatabase:
134
125
  assert TestDatabase.table_base == NoTZBase
135
126
  with Database(
136
127
  ssh_host="mysshhost", ssh_port=25, db_has_tz=True, **params
137
- ) as dbsession:
128
+ ) as dbdatabase:
129
+ dbsession = dbdatabase.get_session()
138
130
  assert (
139
131
  str(dbsession.bind.engine.url)
140
132
  == "postgresql+psycopg://myuser@0.0.0.0:12345/mydatabase"
141
133
  )
142
134
  assert TestDatabase.table_base == TZBase
135
+ with Database(
136
+ ssh_host="mysshhost", ssh_port=25, table_base=TZBase, **params
137
+ ) as dbdatabase:
138
+ dbsession = dbdatabase.get_session()
139
+ assert (
140
+ str(dbsession.bind.engine.url)
141
+ == "postgresql+psycopg://myuser@0.0.0.0:12345/mydatabase"
142
+ )
143
+ assert TestDatabase.table_base == TZBase
144
+
145
+ def test_recreate_schema(self, mock_engine):
146
+ db_uri = "postgresql+psycopg://myuser:mypass@0.0.0.0:12345/mydatabase"
147
+ assert Database.recreate_schema(mock_engine, db_uri) is True
148
+ db_uri = "Error"
149
+ assert Database.recreate_schema(mock_engine, db_uri) is False
@@ -1,2 +0,0 @@
1
- [tool.black]
2
- line-length = 79
@@ -1,26 +0,0 @@
1
- default_language_version:
2
- python: python3.11
3
- repos:
4
- - repo: https://github.com/pre-commit/pre-commit-hooks
5
- rev: v4.4.0
6
- hooks:
7
- - id: trailing-whitespace
8
- - id: end-of-file-fixer
9
- - id: check-ast
10
- - repo: https://github.com/psf/black
11
- rev: 23.3.0
12
- hooks:
13
- - id: black
14
- args: [--config, .config/black.toml]
15
- - repo: https://github.com/charliermarsh/ruff-pre-commit
16
- rev: v0.0.267
17
- hooks:
18
- - id: ruff
19
- args: [--config, .config/ruff.toml, --fix, --exit-non-zero-on-fix]
20
- - repo: https://github.com/jazzband/pip-tools
21
- rev: 6.13.0
22
- hooks:
23
- - id: pip-compile
24
- name: pip-compile requirements.txt
25
- files: pyproject.toml
26
- args: [pyproject.toml, --resolver=backtracking, --all-extras, --upgrade, -q, -o, requirements.txt]
@@ -1,69 +0,0 @@
1
- #
2
- # This file is autogenerated by pip-compile with Python 3.11
3
- # by the following command:
4
- #
5
- # pip-compile --all-extras --output-file=requirements.txt --resolver=backtracking pyproject.toml
6
- #
7
- bcrypt==4.0.1
8
- # via paramiko
9
- cffi==1.16.0
10
- # via
11
- # cryptography
12
- # pynacl
13
- cfgv==3.4.0
14
- # via pre-commit
15
- coverage[toml]==7.3.2
16
- # via pytest-cov
17
- cryptography==41.0.4
18
- # via paramiko
19
- distlib==0.3.7
20
- # via virtualenv
21
- filelock==3.12.4
22
- # via virtualenv
23
- greenlet==3.0.0
24
- # via sqlalchemy
25
- identify==2.5.30
26
- # via pre-commit
27
- iniconfig==2.0.0
28
- # via pytest
29
- nodeenv==1.8.0
30
- # via pre-commit
31
- packaging==23.2
32
- # via pytest
33
- paramiko==3.3.1
34
- # via sshtunnel
35
- platformdirs==3.11.0
36
- # via virtualenv
37
- pluggy==1.3.0
38
- # via pytest
39
- pre-commit==3.4.0
40
- # via hdx-python-database (pyproject.toml)
41
- psycopg[binary]==3.1.12
42
- # via hdx-python-database (pyproject.toml)
43
- psycopg-binary==3.1.12
44
- # via psycopg
45
- pycparser==2.21
46
- # via cffi
47
- pynacl==1.5.0
48
- # via paramiko
49
- pytest==7.4.2
50
- # via
51
- # hdx-python-database (pyproject.toml)
52
- # pytest-cov
53
- pytest-cov==4.1.0
54
- # via hdx-python-database (pyproject.toml)
55
- pyyaml==6.0.1
56
- # via pre-commit
57
- sqlalchemy==2.0.21
58
- # via hdx-python-database (pyproject.toml)
59
- sshtunnel==0.4.0
60
- # via hdx-python-database (pyproject.toml)
61
- typing-extensions==4.8.0
62
- # via
63
- # psycopg
64
- # sqlalchemy
65
- virtualenv==20.24.5
66
- # via pre-commit
67
-
68
- # The following packages are considered to be unsafe in a requirements file:
69
- # setuptools
@@ -1,149 +0,0 @@
1
- """Database utilities"""
2
- import logging
3
- from typing import Any, Optional, Type, Union
4
-
5
- from sqlalchemy import create_engine
6
- from sqlalchemy.ext.automap import automap_base
7
- from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
8
- from sqlalchemy.pool import NullPool
9
-
10
- from ._version import version as __version__ # noqa: F401
11
- from .dburi import get_connection_uri
12
- from .no_timezone import Base as NoTZBase
13
- from .postgresql import wait_for_postgresql
14
- from .with_timezone import Base as TZBase
15
-
16
- logger = logging.getLogger(__name__)
17
-
18
-
19
- class Database:
20
- """Database helper class to handle ssh tunnels, waiting for PostgreSQL to
21
- be up etc. Can be used in a with statement returning a Session object that
22
- if reflect is True will have a variable reflected_classes containing the
23
- reflected classes. db_has_tz which defaults to False indicates whether
24
- database datetime columns have timezones. If not, conversion occurs between
25
- Python datetimes with timezones to timezoneless database columns (but not
26
- when using reflection).
27
-
28
- Args:
29
- database (Optional[str]): Database name
30
- host (Optional[str]): Host where database is located
31
- port (Union[int, str, None]): Database port
32
- username (Optional[str]): Username to log into database
33
- password (Optional[str]): Password to log into database
34
- dialect (str): Database dialect. Defaults to "postgresql".
35
- driver (Optional[str]): Database driver. Defaults to None (psycopg if postgresql or None)
36
- db_has_tz (bool): True if db datetime columns have timezone. Defaults to False.
37
- reflect (bool): Whether to reflect existing tables. Defaults to False.
38
- **kwargs: See below
39
- ssh_host (str): SSH host (the server to connect to)
40
- ssh_port (int): SSH port. Defaults to 22.
41
- ssh_username (str): SSH username
42
- ssh_password (str): SSH password
43
- ssh_private_key: Location of SSH private key (instead of ssh_password)
44
- For more advanced usage, see SSHTunnelForwarder documentation.
45
-
46
- """
47
-
48
- def __init__(
49
- self,
50
- database: Optional[str] = None,
51
- host: Optional[str] = None,
52
- port: Union[int, str, None] = None,
53
- username: Optional[str] = None,
54
- password: Optional[str] = None,
55
- dialect: str = "postgresql",
56
- driver: Optional[str] = None,
57
- db_has_tz: bool = False,
58
- reflect: bool = False,
59
- **kwargs: Any,
60
- ) -> None:
61
- if port is not None:
62
- port = int(port)
63
- if len(kwargs) != 0:
64
- try:
65
- import sshtunnel
66
- except ImportError:
67
- # dependency missing, log an error
68
- logger.error(
69
- "sshtunnel not found! Please install hdx-python-database[sshtunnel] to enable."
70
- )
71
- raise
72
- ssh_host = kwargs["ssh_host"]
73
- del kwargs["ssh_host"]
74
- ssh_port = kwargs.get("ssh_port")
75
- if ssh_port is not None:
76
- ssh_port = int(ssh_port)
77
- del kwargs["ssh_port"]
78
- else:
79
- ssh_port = 22
80
- self.server = sshtunnel.SSHTunnelForwarder(
81
- (ssh_host, ssh_port),
82
- remote_bind_address=(host, port),
83
- **kwargs,
84
- )
85
- self.server.start()
86
- host = self.server.local_bind_host
87
- port = self.server.local_bind_port
88
- else:
89
- self.server = None
90
- db_uri = get_connection_uri(
91
- database,
92
- host,
93
- port,
94
- username,
95
- password,
96
- dialect=dialect,
97
- driver=driver,
98
- )
99
- if dialect == "postgresql":
100
- wait_for_postgresql(db_uri)
101
- if db_has_tz:
102
- table_base = TZBase
103
- else:
104
- table_base = NoTZBase
105
- self.session = self.get_session(
106
- db_uri, table_base=table_base, reflect=reflect
107
- )
108
-
109
- def __enter__(self) -> Session:
110
- return self.session
111
-
112
- def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
113
- self.session.close()
114
- if self.server is not None:
115
- self.server.stop()
116
-
117
- @staticmethod
118
- def get_session(
119
- db_uri: str,
120
- table_base: Type[DeclarativeBase] = NoTZBase,
121
- reflect: bool = False,
122
- ) -> Session:
123
- """Gets SQLAlchemy session given url. Tables must inherit from Base in
124
- hdx.utilities.database unless base is defined. If reflect is True,
125
- classes will be reflected from an existing database and the reflected
126
- classes are returned in a variable reflected_classes in the returned
127
- Session object. Note that type annotation maps don't work with
128
- reflection.
129
-
130
- Args:
131
- db_uri (str): Connection URI
132
- table_base (Type[DeclarativeBase]): Base database table class. Defaults to NoTZBase.
133
- reflect (bool): Whether to reflect existing tables. Defaults to False.
134
-
135
- Returns:
136
- sqlalchemy.orm.Session: SQLAlchemy session
137
- """
138
- engine = create_engine(db_uri, poolclass=NullPool, echo=False)
139
- Session = sessionmaker(bind=engine)
140
- if reflect:
141
- Base = automap_base(declarative_base=table_base)
142
- Base.prepare(autoload_with=engine)
143
- classes = Base.classes
144
- else:
145
- table_base.metadata.create_all(engine)
146
- classes = None
147
- session = Session()
148
- session.reflected_classes = classes
149
- return session
@@ -1,22 +0,0 @@
1
- import psycopg
2
- import pytest
3
-
4
- from . import PsycopgConnection
5
-
6
-
7
- @pytest.fixture(scope="function")
8
- def mock_psycopg(monkeypatch):
9
- def connect(*args, **kwargs):
10
- if PsycopgConnection.connected:
11
-
12
- class Connection:
13
- @staticmethod
14
- def close():
15
- return
16
-
17
- return Connection()
18
- else:
19
- PsycopgConnection.connected = True
20
- raise psycopg.OperationalError
21
-
22
- monkeypatch.setattr(psycopg, "connect", connect)