We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Leaking DB connections in intergration tests
Integration tests started occasionally failing with the following error:

```
psycopg2.OperationalError: FATAL:  sorry, too many clients already
```

I have not been able to reproduce this in the local dev environment, but it happens on Github Actions from time to time, typically when running ` integration_tests/test_dataset_add.py` suite.

Integration tests run  `datacube` CLI app, but in a fancy way, avoiding launching a separate process. This is needed for accurate coverage tracking, but I feel like this leads to resource leakage. Apps are designed for once off running, so they might not care about closing db connections in a "proper way".

Acceptable solutions:

- Bump connection pool to avoid problems
- Find a way for test fixtures to cleanup properly



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `datacube/ui/click.py`
Content:
```
1 """
2 Common functions for click-based cli scripts.
3 """
4 import functools
5 import logging
6 import os
7 import copy
8 import sys
9 
10 import click
11 
12 from datacube import config, __version__
13 from datacube.api.core import Datacube
14 
15 from datacube.executor import get_executor, mk_celery_executor
16 from datacube.index import index_connect
17 from pathlib import Path
18 
19 from datacube.ui.expression import parse_expressions
20 from sqlalchemy.exc import OperationalError, ProgrammingError
21 
22 _LOG_FORMAT_STRING = '%(asctime)s %(process)d %(name)s %(levelname)s %(message)s'
23 CLICK_SETTINGS = dict(help_option_names=['-h', '--help'])
24 _LOG = logging.getLogger(__name__)
25 
26 
27 def _print_version(ctx, param, value):
28     if not value or ctx.resilient_parsing:
29         return
30 
31     click.echo(
32         '{prog}, version {version}'.format(
33             prog='Open Data Cube core',
34             version=__version__
35         )
36     )
37     ctx.exit()
38 
39 
40 def compose(*functions):
41     """
42     >>> compose(
43     ...     lambda x: x+1,
44     ...     lambda y: y+2
45     ... )(1)
46     4
47     """
48 
49     def compose2(f, g):
50         return lambda x: f(g(x))
51 
52     return functools.reduce(compose2, functions, lambda x: x)
53 
54 
55 class ColorFormatter(logging.Formatter):
56     colors = {
57         'info': dict(fg='white'),
58         'error': dict(fg='red'),
59         'exception': dict(fg='red'),
60         'critical': dict(fg='red'),
61         'debug': dict(fg='blue'),
62         'warning': dict(fg='yellow')
63     }
64 
65     def format(self, record):
66         if not record.exc_info:
67             record = copy.copy(record)
68             record.levelname = click.style(record.levelname, **self.colors.get(record.levelname.lower(), {}))
69         return logging.Formatter.format(self, record)
70 
71 
72 class ClickHandler(logging.Handler):
73     def emit(self, record):
74         try:
75             msg = self.format(record)
76             click.echo(msg, err=True)
77         except:  # pylint: disable=bare-except
78             self.handleError(record)
79 
80 
81 def remove_handlers_of_type(logger, handler_type):
82     for handler in logger.handlers:
83         if isinstance(handler, handler_type):
84             logger.removeHandler(handler)
85 
86 
87 def _init_logging(ctx, param, value):
88     # When running in tests, we don't want to keep adding log handlers. It creates duplicate log messages up the wahoo.
89     remove_handlers_of_type(logging.root, ClickHandler)
90     handler = ClickHandler()
91     handler.formatter = ColorFormatter(_LOG_FORMAT_STRING)
92     logging.root.addHandler(handler)
93 
94     logging_level = logging.WARN - 10 * value
95     logging.root.setLevel(logging_level)
96     logging.getLogger('datacube').setLevel(logging_level)
97 
98     if logging_level <= logging.INFO:
99         logging.getLogger('rasterio').setLevel(logging.INFO)
100 
101     logging.getLogger('datacube').info('Running datacube command: %s', ' '.join(sys.argv))
102 
103     if not ctx.obj:
104         ctx.obj = {}
105 
106     ctx.obj['verbosity'] = value
107 
108 
109 def _add_logfile(ctx, param, value):
110     formatter = logging.Formatter(_LOG_FORMAT_STRING)
111     for logfile in value:
112         handler = logging.FileHandler(logfile)
113         handler.formatter = formatter
114         logging.root.addHandler(handler)
115 
116 
117 def _log_queries(ctx, param, value):
118     if value:
119         logging.getLogger('sqlalchemy.engine').setLevel('INFO')
120 
121 
122 def _set_config(ctx, param, value):
123     if value:
124         if not any(os.path.exists(p) for p in value):
125             raise ValueError('No specified config paths exist: {}'.format(value))
126 
127         if not ctx.obj:
128             ctx.obj = {}
129         paths = value
130         ctx.obj['config_files'] = paths
131     return value
132 
133 
134 def _set_environment(ctx, param, value):
135     if not ctx.obj:
136         ctx.obj = {}
137     ctx.obj['config_environment'] = value
138 
139 
140 #: pylint: disable=invalid-name
141 version_option = click.option('--version', is_flag=True, callback=_print_version,
142                               expose_value=False, is_eager=True)
143 #: pylint: disable=invalid-name
144 verbose_option = click.option('--verbose', '-v', count=True, callback=_init_logging,
145                               is_eager=True, expose_value=False, help="Use multiple times for more verbosity")
146 #: pylint: disable=invalid-name
147 logfile_option = click.option('--log-file', multiple=True, callback=_add_logfile,
148                               is_eager=True, expose_value=False, help="Specify log file")
149 #: pylint: disable=invalid-name
150 config_option = click.option('--config', '--config_file', '-C', multiple=True, default='', callback=_set_config,
151                              expose_value=False)
152 config_option_exposed = click.option('--config', '--config_file', '-C', multiple=True, default='', callback=_set_config)
153 
154 environment_option = click.option('--env', '-E', callback=_set_environment,
155                                   expose_value=False)
156 #: pylint: disable=invalid-name
157 log_queries_option = click.option('--log-queries', is_flag=True, callback=_log_queries,
158                                   expose_value=False, help="Print database queries.")
159 
160 # This is a function, so it's valid to be lowercase.
161 #: pylint: disable=invalid-name
162 global_cli_options = compose(
163     version_option,
164     verbose_option,
165     logfile_option,
166     environment_option,
167     config_option,
168     log_queries_option
169 )
170 
171 
172 @click.group(help="Data Cube command-line interface", context_settings=CLICK_SETTINGS)
173 @global_cli_options
174 def cli():
175     pass
176 
177 
178 def pass_config(f):
179     """Get a datacube config as the first argument. """
180 
181     def new_func(*args, **kwargs):
182         obj = click.get_current_context().obj
183 
184         paths = obj.get('config_files', None)
185         # If the user is overriding the defaults
186         specific_environment = obj.get('config_environment')
187 
188         try:
189             parsed_config = config.LocalConfig.find(paths=paths, env=specific_environment)
190         except ValueError:
191             if specific_environment:
192                 raise click.ClickException("No datacube config found for '{}'".format(specific_environment))
193             else:
194                 raise click.ClickException("No datacube config found")
195 
196         _LOG.debug("Loaded datacube config: %r", parsed_config)
197         return f(parsed_config, *args, **kwargs)
198 
199     return functools.update_wrapper(new_func, f)
200 
201 
202 def pass_index(app_name=None, expect_initialised=True):
203     """Get a connection to the index as the first argument.
204 
205     :param str app_name:
206         A short name of the application for logging purposes.
207     :param bool expect_initialised:
208         Whether to connect immediately on startup. Useful to catch connection config issues immediately,
209         but if you're planning to fork before any usage (such as in the case of some web servers),
210         you may not want this. For more information on thread/process usage, see datacube.index.Index
211     """
212 
213     def decorate(f):
214         @pass_config
215         def with_index(local_config: config.LocalConfig,
216                        *args,
217                        **kwargs):
218             ctx = click.get_current_context()
219             try:
220                 index = index_connect(local_config,
221                                       application_name=app_name or ctx.command_path,
222                                       validate_connection=expect_initialised)
223                 ctx.obj['index'] = index
224                 _LOG.debug("Connected to datacube index: %s", index)
225                 return f(index, *args, **kwargs)
226             except (OperationalError, ProgrammingError) as e:
227                 handle_exception('Error Connecting to database: %s', e)
228 
229         return functools.update_wrapper(with_index, f)
230 
231     return decorate
232 
233 
234 def pass_datacube(app_name=None, expect_initialised=True):
235     """
236     Get a DataCube from the current or default local settings.
237 
238     :param str app_name:
239         A short name of the application for logging purposes.
240     :param bool expect_initialised:
241         Whether to connect immediately on startup. Useful to catch connection config issues immediately,
242         but if you're planning to fork before any usage (such as in the case of some web servers),
243         you may not want this. For More information on thread/process usage, see datacube.index.Index
244     """
245 
246     def decorate(f):
247         @pass_index(app_name=app_name, expect_initialised=expect_initialised)
248         def with_datacube(index, *args, **kwargs):
249             return f(Datacube(index=index), *args, **kwargs)
250 
251         return functools.update_wrapper(with_datacube, f)
252 
253     return decorate
254 
255 
256 def parse_endpoint(value):
257     ip, port = tuple(value.split(':'))
258     return ip, int(port)
259 
260 
261 EXECUTOR_TYPES = {
262     'serial': lambda _: get_executor(None, None),
263     'multiproc': lambda workers: get_executor(None, int(workers)),
264     'distributed': lambda addr: get_executor(parse_endpoint(addr), True),
265     'celery': lambda addr: mk_celery_executor(*parse_endpoint(addr))
266 }
267 
268 EXECUTOR_TYPES['dask'] = EXECUTOR_TYPES['distributed']  # Add alias "dask" for distributed
269 
270 
271 def _setup_executor(ctx, param, value):
272     try:
273         return EXECUTOR_TYPES[value[0]](value[1])
274     except ValueError:
275         ctx.fail("Failed to create '%s' executor with '%s'" % value)
276 
277 
278 executor_cli_options = click.option('--executor',  # type: ignore
279                                     type=(click.Choice(list(EXECUTOR_TYPES)), str),
280                                     default=('serial', None),
281                                     help="Run parallelized, either locally or distributed. eg:\n"
282                                          "--executor multiproc 4 (OR)\n"
283                                          "--executor distributed 10.0.0.8:8888",
284                                     callback=_setup_executor)
285 
286 
287 def handle_exception(msg, e):
288     """
289     Exit following an exception in a CLI app
290 
291     If verbosity (-v flag) specified, dump out a stack trace. Otherwise,
292     simply print the given error message.
293 
294     Include a '%s' in the message to print the single line message from the
295     exception.
296 
297     :param e: caught Exception
298     :param msg: Message to User with optional %s
299     """
300     ctx = click.get_current_context()
301     if ctx.obj['verbosity'] >= 1:
302         raise e
303     else:
304         if '%s' in msg:
305             click.echo(msg % e)
306         else:
307             click.echo(msg)
308         ctx.exit(1)
309 
310 
311 def parsed_search_expressions(f):
312     """
313     Add [EXPRESSIONs] arguments to a click application
314 
315     Passes a parsed dict of search expressions to the `expressions` argument
316     of the command.
317 
318     Also appends documentation on using search expressions to the command.
319 
320     WARNING: This wrapped expects an unlimited number of search expressions
321     as click arguments, which means your command must take only click options
322     or a specified number of arguments.
323     """
324     if not f.__doc__:
325         f.__doc__ = ""
326     f.__doc__ += """
327     EXPRESSIONS
328     
329     Select datasets using [EXPRESSIONS] to filter by date, product type,
330     spatial extents or other searchable fields.
331 
332     \b
333         FIELD = VALUE
334         FIELD in DATE-RANGE
335         FIELD in [START, END]
336 
337     \b
338     DATE-RANGE is one of YYYY, YYYY-MM or YYYY-MM-DD
339     START and END can be either numbers or dates
340 
341     FIELD: x, y, lat, lon, time, product, ... 
342 
343     \b
344     eg. 'time in [1996-01-01, 1996-12-31]'
345         'time in 1996'
346         'lon in [130, 140]' 'lat in [-40, -30]'
347         product=ls5_nbar_albers
348 
349     """
350 
351     def my_parse(ctx, param, value):
352         return parse_expressions(*list(value))
353 
354     f = click.argument('expressions', callback=my_parse, nargs=-1)(f)
355     return f
356
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/datacube/ui/click.py b/datacube/ui/click.py
--- a/datacube/ui/click.py
+++ b/datacube/ui/click.py
@@ -215,16 +215,21 @@
         def with_index(local_config: config.LocalConfig,
                        *args,
                        **kwargs):
-            ctx = click.get_current_context()
+            command_path = click.get_current_context().command_path
             try:
                 index = index_connect(local_config,
-                                      application_name=app_name or ctx.command_path,
+                                      application_name=app_name or command_path,
                                       validate_connection=expect_initialised)
-                ctx.obj['index'] = index
                 _LOG.debug("Connected to datacube index: %s", index)
-                return f(index, *args, **kwargs)
             except (OperationalError, ProgrammingError) as e:
                 handle_exception('Error Connecting to database: %s', e)
+                return
+
+            try:
+                return f(index, *args, **kwargs)
+            finally:
+                index.close()
+                del index
 
         return functools.update_wrapper(with_index, f)

{"golden_diff": "diff --git a/datacube/ui/click.py b/datacube/ui/click.py\n--- a/datacube/ui/click.py\n+++ b/datacube/ui/click.py\n@@ -215,16 +215,21 @@\n         def with_index(local_config: config.LocalConfig,\n                        *args,\n                        **kwargs):\n-            ctx = click.get_current_context()\n+            command_path = click.get_current_context().command_path\n             try:\n                 index = index_connect(local_config,\n-                                      application_name=app_name or ctx.command_path,\n+                                      application_name=app_name or command_path,\n                                       validate_connection=expect_initialised)\n-                ctx.obj['index'] = index\n                 _LOG.debug(\"Connected to datacube index: %s\", index)\n-                return f(index, *args, **kwargs)\n             except (OperationalError, ProgrammingError) as e:\n                 handle_exception('Error Connecting to database: %s', e)\n+                return\n+\n+            try:\n+                return f(index, *args, **kwargs)\n+            finally:\n+                index.close()\n+                del index\n \n         return functools.update_wrapper(with_index, f)\n", "issue": "Leaking DB connections in intergration tests\nIntegration tests started occasionally failing with the following error:\r\n\r\n```\r\npsycopg2.OperationalError: FATAL:  sorry, too many clients already\r\n```\r\n\r\nI have not been able to reproduce this in the local dev environment, but it happens on Github Actions from time to time, typically when running ` integration_tests/test_dataset_add.py` suite.\r\n\r\nIntegration tests run  `datacube` CLI app, but in a fancy way, avoiding launching a separate process. This is needed for accurate coverage tracking, but I feel like this leads to resource leakage. Apps are designed for once off running, so they might not care about closing db connections in a \"proper way\".\r\n\r\nAcceptable solutions:\r\n\r\n- Bump connection pool to avoid problems\r\n- Find a way for test fixtures to cleanup properly\r\n\r\n\n", "before_files": [{"content": "\"\"\"\nCommon functions for click-based cli scripts.\n\"\"\"\nimport functools\nimport logging\nimport os\nimport copy\nimport sys\n\nimport click\n\nfrom datacube import config, __version__\nfrom datacube.api.core import Datacube\n\nfrom datacube.executor import get_executor, mk_celery_executor\nfrom datacube.index import index_connect\nfrom pathlib import Path\n\nfrom datacube.ui.expression import parse_expressions\nfrom sqlalchemy.exc import OperationalError, ProgrammingError\n\n_LOG_FORMAT_STRING = '%(asctime)s %(process)d %(name)s %(levelname)s %(message)s'\nCLICK_SETTINGS = dict(help_option_names=['-h', '--help'])\n_LOG = logging.getLogger(__name__)\n\n\ndef _print_version(ctx, param, value):\n    if not value or ctx.resilient_parsing:\n        return\n\n    click.echo(\n        '{prog}, version {version}'.format(\n            prog='Open Data Cube core',\n            version=__version__\n        )\n    )\n    ctx.exit()\n\n\ndef compose(*functions):\n    \"\"\"\n    >>> compose(\n    ...     lambda x: x+1,\n    ...     lambda y: y+2\n    ... )(1)\n    4\n    \"\"\"\n\n    def compose2(f, g):\n        return lambda x: f(g(x))\n\n    return functools.reduce(compose2, functions, lambda x: x)\n\n\nclass ColorFormatter(logging.Formatter):\n    colors = {\n        'info': dict(fg='white'),\n        'error': dict(fg='red'),\n        'exception': dict(fg='red'),\n        'critical': dict(fg='red'),\n        'debug': dict(fg='blue'),\n        'warning': dict(fg='yellow')\n    }\n\n    def format(self, record):\n        if not record.exc_info:\n            record = copy.copy(record)\n            record.levelname = click.style(record.levelname, **self.colors.get(record.levelname.lower(), {}))\n        return logging.Formatter.format(self, record)\n\n\nclass ClickHandler(logging.Handler):\n    def emit(self, record):\n        try:\n            msg = self.format(record)\n            click.echo(msg, err=True)\n        except:  # pylint: disable=bare-except\n            self.handleError(record)\n\n\ndef remove_handlers_of_type(logger, handler_type):\n    for handler in logger.handlers:\n        if isinstance(handler, handler_type):\n            logger.removeHandler(handler)\n\n\ndef _init_logging(ctx, param, value):\n    # When running in tests, we don't want to keep adding log handlers. It creates duplicate log messages up the wahoo.\n    remove_handlers_of_type(logging.root, ClickHandler)\n    handler = ClickHandler()\n    handler.formatter = ColorFormatter(_LOG_FORMAT_STRING)\n    logging.root.addHandler(handler)\n\n    logging_level = logging.WARN - 10 * value\n    logging.root.setLevel(logging_level)\n    logging.getLogger('datacube').setLevel(logging_level)\n\n    if logging_level <= logging.INFO:\n        logging.getLogger('rasterio').setLevel(logging.INFO)\n\n    logging.getLogger('datacube').info('Running datacube command: %s', ' '.join(sys.argv))\n\n    if not ctx.obj:\n        ctx.obj = {}\n\n    ctx.obj['verbosity'] = value\n\n\ndef _add_logfile(ctx, param, value):\n    formatter = logging.Formatter(_LOG_FORMAT_STRING)\n    for logfile in value:\n        handler = logging.FileHandler(logfile)\n        handler.formatter = formatter\n        logging.root.addHandler(handler)\n\n\ndef _log_queries(ctx, param, value):\n    if value:\n        logging.getLogger('sqlalchemy.engine').setLevel('INFO')\n\n\ndef _set_config(ctx, param, value):\n    if value:\n        if not any(os.path.exists(p) for p in value):\n            raise ValueError('No specified config paths exist: {}'.format(value))\n\n        if not ctx.obj:\n            ctx.obj = {}\n        paths = value\n        ctx.obj['config_files'] = paths\n    return value\n\n\ndef _set_environment(ctx, param, value):\n    if not ctx.obj:\n        ctx.obj = {}\n    ctx.obj['config_environment'] = value\n\n\n#: pylint: disable=invalid-name\nversion_option = click.option('--version', is_flag=True, callback=_print_version,\n                              expose_value=False, is_eager=True)\n#: pylint: disable=invalid-name\nverbose_option = click.option('--verbose', '-v', count=True, callback=_init_logging,\n                              is_eager=True, expose_value=False, help=\"Use multiple times for more verbosity\")\n#: pylint: disable=invalid-name\nlogfile_option = click.option('--log-file', multiple=True, callback=_add_logfile,\n                              is_eager=True, expose_value=False, help=\"Specify log file\")\n#: pylint: disable=invalid-name\nconfig_option = click.option('--config', '--config_file', '-C', multiple=True, default='', callback=_set_config,\n                             expose_value=False)\nconfig_option_exposed = click.option('--config', '--config_file', '-C', multiple=True, default='', callback=_set_config)\n\nenvironment_option = click.option('--env', '-E', callback=_set_environment,\n                                  expose_value=False)\n#: pylint: disable=invalid-name\nlog_queries_option = click.option('--log-queries', is_flag=True, callback=_log_queries,\n                                  expose_value=False, help=\"Print database queries.\")\n\n# This is a function, so it's valid to be lowercase.\n#: pylint: disable=invalid-name\nglobal_cli_options = compose(\n    version_option,\n    verbose_option,\n    logfile_option,\n    environment_option,\n    config_option,\n    log_queries_option\n)\n\n\n@click.group(help=\"Data Cube command-line interface\", context_settings=CLICK_SETTINGS)\n@global_cli_options\ndef cli():\n    pass\n\n\ndef pass_config(f):\n    \"\"\"Get a datacube config as the first argument. \"\"\"\n\n    def new_func(*args, **kwargs):\n        obj = click.get_current_context().obj\n\n        paths = obj.get('config_files', None)\n        # If the user is overriding the defaults\n        specific_environment = obj.get('config_environment')\n\n        try:\n            parsed_config = config.LocalConfig.find(paths=paths, env=specific_environment)\n        except ValueError:\n            if specific_environment:\n                raise click.ClickException(\"No datacube config found for '{}'\".format(specific_environment))\n            else:\n                raise click.ClickException(\"No datacube config found\")\n\n        _LOG.debug(\"Loaded datacube config: %r\", parsed_config)\n        return f(parsed_config, *args, **kwargs)\n\n    return functools.update_wrapper(new_func, f)\n\n\ndef pass_index(app_name=None, expect_initialised=True):\n    \"\"\"Get a connection to the index as the first argument.\n\n    :param str app_name:\n        A short name of the application for logging purposes.\n    :param bool expect_initialised:\n        Whether to connect immediately on startup. Useful to catch connection config issues immediately,\n        but if you're planning to fork before any usage (such as in the case of some web servers),\n        you may not want this. For more information on thread/process usage, see datacube.index.Index\n    \"\"\"\n\n    def decorate(f):\n        @pass_config\n        def with_index(local_config: config.LocalConfig,\n                       *args,\n                       **kwargs):\n            ctx = click.get_current_context()\n            try:\n                index = index_connect(local_config,\n                                      application_name=app_name or ctx.command_path,\n                                      validate_connection=expect_initialised)\n                ctx.obj['index'] = index\n                _LOG.debug(\"Connected to datacube index: %s\", index)\n                return f(index, *args, **kwargs)\n            except (OperationalError, ProgrammingError) as e:\n                handle_exception('Error Connecting to database: %s', e)\n\n        return functools.update_wrapper(with_index, f)\n\n    return decorate\n\n\ndef pass_datacube(app_name=None, expect_initialised=True):\n    \"\"\"\n    Get a DataCube from the current or default local settings.\n\n    :param str app_name:\n        A short name of the application for logging purposes.\n    :param bool expect_initialised:\n        Whether to connect immediately on startup. Useful to catch connection config issues immediately,\n        but if you're planning to fork before any usage (such as in the case of some web servers),\n        you may not want this. For More information on thread/process usage, see datacube.index.Index\n    \"\"\"\n\n    def decorate(f):\n        @pass_index(app_name=app_name, expect_initialised=expect_initialised)\n        def with_datacube(index, *args, **kwargs):\n            return f(Datacube(index=index), *args, **kwargs)\n\n        return functools.update_wrapper(with_datacube, f)\n\n    return decorate\n\n\ndef parse_endpoint(value):\n    ip, port = tuple(value.split(':'))\n    return ip, int(port)\n\n\nEXECUTOR_TYPES = {\n    'serial': lambda _: get_executor(None, None),\n    'multiproc': lambda workers: get_executor(None, int(workers)),\n    'distributed': lambda addr: get_executor(parse_endpoint(addr), True),\n    'celery': lambda addr: mk_celery_executor(*parse_endpoint(addr))\n}\n\nEXECUTOR_TYPES['dask'] = EXECUTOR_TYPES['distributed']  # Add alias \"dask\" for distributed\n\n\ndef _setup_executor(ctx, param, value):\n    try:\n        return EXECUTOR_TYPES[value[0]](value[1])\n    except ValueError:\n        ctx.fail(\"Failed to create '%s' executor with '%s'\" % value)\n\n\nexecutor_cli_options = click.option('--executor',  # type: ignore\n                                    type=(click.Choice(list(EXECUTOR_TYPES)), str),\n                                    default=('serial', None),\n                                    help=\"Run parallelized, either locally or distributed. eg:\\n\"\n                                         \"--executor multiproc 4 (OR)\\n\"\n                                         \"--executor distributed 10.0.0.8:8888\",\n                                    callback=_setup_executor)\n\n\ndef handle_exception(msg, e):\n    \"\"\"\n    Exit following an exception in a CLI app\n\n    If verbosity (-v flag) specified, dump out a stack trace. Otherwise,\n    simply print the given error message.\n\n    Include a '%s' in the message to print the single line message from the\n    exception.\n\n    :param e: caught Exception\n    :param msg: Message to User with optional %s\n    \"\"\"\n    ctx = click.get_current_context()\n    if ctx.obj['verbosity'] >= 1:\n        raise e\n    else:\n        if '%s' in msg:\n            click.echo(msg % e)\n        else:\n            click.echo(msg)\n        ctx.exit(1)\n\n\ndef parsed_search_expressions(f):\n    \"\"\"\n    Add [EXPRESSIONs] arguments to a click application\n\n    Passes a parsed dict of search expressions to the `expressions` argument\n    of the command.\n\n    Also appends documentation on using search expressions to the command.\n\n    WARNING: This wrapped expects an unlimited number of search expressions\n    as click arguments, which means your command must take only click options\n    or a specified number of arguments.\n    \"\"\"\n    if not f.__doc__:\n        f.__doc__ = \"\"\n    f.__doc__ += \"\"\"\n    EXPRESSIONS\n    \n    Select datasets using [EXPRESSIONS] to filter by date, product type,\n    spatial extents or other searchable fields.\n\n    \\b\n        FIELD = VALUE\n        FIELD in DATE-RANGE\n        FIELD in [START, END]\n\n    \\b\n    DATE-RANGE is one of YYYY, YYYY-MM or YYYY-MM-DD\n    START and END can be either numbers or dates\n\n    FIELD: x, y, lat, lon, time, product, ... \n\n    \\b\n    eg. 'time in [1996-01-01, 1996-12-31]'\n        'time in 1996'\n        'lon in [130, 140]' 'lat in [-40, -30]'\n        product=ls5_nbar_albers\n\n    \"\"\"\n\n    def my_parse(ctx, param, value):\n        return parse_expressions(*list(value))\n\n    f = click.argument('expressions', callback=my_parse, nargs=-1)(f)\n    return f\n", "path": "datacube/ui/click.py"}], "after_files": [{"content": "\"\"\"\nCommon functions for click-based cli scripts.\n\"\"\"\nimport functools\nimport logging\nimport os\nimport copy\nimport sys\n\nimport click\n\nfrom datacube import config, __version__\nfrom datacube.api.core import Datacube\n\nfrom datacube.executor import get_executor, mk_celery_executor\nfrom datacube.index import index_connect\nfrom pathlib import Path\n\nfrom datacube.ui.expression import parse_expressions\nfrom sqlalchemy.exc import OperationalError, ProgrammingError\n\n_LOG_FORMAT_STRING = '%(asctime)s %(process)d %(name)s %(levelname)s %(message)s'\nCLICK_SETTINGS = dict(help_option_names=['-h', '--help'])\n_LOG = logging.getLogger(__name__)\n\n\ndef _print_version(ctx, param, value):\n    if not value or ctx.resilient_parsing:\n        return\n\n    click.echo(\n        '{prog}, version {version}'.format(\n            prog='Open Data Cube core',\n            version=__version__\n        )\n    )\n    ctx.exit()\n\n\ndef compose(*functions):\n    \"\"\"\n    >>> compose(\n    ...     lambda x: x+1,\n    ...     lambda y: y+2\n    ... )(1)\n    4\n    \"\"\"\n\n    def compose2(f, g):\n        return lambda x: f(g(x))\n\n    return functools.reduce(compose2, functions, lambda x: x)\n\n\nclass ColorFormatter(logging.Formatter):\n    colors = {\n        'info': dict(fg='white'),\n        'error': dict(fg='red'),\n        'exception': dict(fg='red'),\n        'critical': dict(fg='red'),\n        'debug': dict(fg='blue'),\n        'warning': dict(fg='yellow')\n    }\n\n    def format(self, record):\n        if not record.exc_info:\n            record = copy.copy(record)\n            record.levelname = click.style(record.levelname, **self.colors.get(record.levelname.lower(), {}))\n        return logging.Formatter.format(self, record)\n\n\nclass ClickHandler(logging.Handler):\n    def emit(self, record):\n        try:\n            msg = self.format(record)\n            click.echo(msg, err=True)\n        except:  # pylint: disable=bare-except\n            self.handleError(record)\n\n\ndef remove_handlers_of_type(logger, handler_type):\n    for handler in logger.handlers:\n        if isinstance(handler, handler_type):\n            logger.removeHandler(handler)\n\n\ndef _init_logging(ctx, param, value):\n    # When running in tests, we don't want to keep adding log handlers. It creates duplicate log messages up the wahoo.\n    remove_handlers_of_type(logging.root, ClickHandler)\n    handler = ClickHandler()\n    handler.formatter = ColorFormatter(_LOG_FORMAT_STRING)\n    logging.root.addHandler(handler)\n\n    logging_level = logging.WARN - 10 * value\n    logging.root.setLevel(logging_level)\n    logging.getLogger('datacube').setLevel(logging_level)\n\n    if logging_level <= logging.INFO:\n        logging.getLogger('rasterio').setLevel(logging.INFO)\n\n    logging.getLogger('datacube').info('Running datacube command: %s', ' '.join(sys.argv))\n\n    if not ctx.obj:\n        ctx.obj = {}\n\n    ctx.obj['verbosity'] = value\n\n\ndef _add_logfile(ctx, param, value):\n    formatter = logging.Formatter(_LOG_FORMAT_STRING)\n    for logfile in value:\n        handler = logging.FileHandler(logfile)\n        handler.formatter = formatter\n        logging.root.addHandler(handler)\n\n\ndef _log_queries(ctx, param, value):\n    if value:\n        logging.getLogger('sqlalchemy.engine').setLevel('INFO')\n\n\ndef _set_config(ctx, param, value):\n    if value:\n        if not any(os.path.exists(p) for p in value):\n            raise ValueError('No specified config paths exist: {}'.format(value))\n\n        if not ctx.obj:\n            ctx.obj = {}\n        paths = value\n        ctx.obj['config_files'] = paths\n    return value\n\n\ndef _set_environment(ctx, param, value):\n    if not ctx.obj:\n        ctx.obj = {}\n    ctx.obj['config_environment'] = value\n\n\n#: pylint: disable=invalid-name\nversion_option = click.option('--version', is_flag=True, callback=_print_version,\n                              expose_value=False, is_eager=True)\n#: pylint: disable=invalid-name\nverbose_option = click.option('--verbose', '-v', count=True, callback=_init_logging,\n                              is_eager=True, expose_value=False, help=\"Use multiple times for more verbosity\")\n#: pylint: disable=invalid-name\nlogfile_option = click.option('--log-file', multiple=True, callback=_add_logfile,\n                              is_eager=True, expose_value=False, help=\"Specify log file\")\n#: pylint: disable=invalid-name\nconfig_option = click.option('--config', '--config_file', '-C', multiple=True, default='', callback=_set_config,\n                             expose_value=False)\nconfig_option_exposed = click.option('--config', '--config_file', '-C', multiple=True, default='', callback=_set_config)\n\nenvironment_option = click.option('--env', '-E', callback=_set_environment,\n                                  expose_value=False)\n#: pylint: disable=invalid-name\nlog_queries_option = click.option('--log-queries', is_flag=True, callback=_log_queries,\n                                  expose_value=False, help=\"Print database queries.\")\n\n# This is a function, so it's valid to be lowercase.\n#: pylint: disable=invalid-name\nglobal_cli_options = compose(\n    version_option,\n    verbose_option,\n    logfile_option,\n    environment_option,\n    config_option,\n    log_queries_option\n)\n\n\n@click.group(help=\"Data Cube command-line interface\", context_settings=CLICK_SETTINGS)\n@global_cli_options\ndef cli():\n    pass\n\n\ndef pass_config(f):\n    \"\"\"Get a datacube config as the first argument. \"\"\"\n\n    def new_func(*args, **kwargs):\n        obj = click.get_current_context().obj\n\n        paths = obj.get('config_files', None)\n        # If the user is overriding the defaults\n        specific_environment = obj.get('config_environment')\n\n        try:\n            parsed_config = config.LocalConfig.find(paths=paths, env=specific_environment)\n        except ValueError:\n            if specific_environment:\n                raise click.ClickException(\"No datacube config found for '{}'\".format(specific_environment))\n            else:\n                raise click.ClickException(\"No datacube config found\")\n\n        _LOG.debug(\"Loaded datacube config: %r\", parsed_config)\n        return f(parsed_config, *args, **kwargs)\n\n    return functools.update_wrapper(new_func, f)\n\n\ndef pass_index(app_name=None, expect_initialised=True):\n    \"\"\"Get a connection to the index as the first argument.\n\n    :param str app_name:\n        A short name of the application for logging purposes.\n    :param bool expect_initialised:\n        Whether to connect immediately on startup. Useful to catch connection config issues immediately,\n        but if you're planning to fork before any usage (such as in the case of some web servers),\n        you may not want this. For more information on thread/process usage, see datacube.index.Index\n    \"\"\"\n\n    def decorate(f):\n        @pass_config\n        def with_index(local_config: config.LocalConfig,\n                       *args,\n                       **kwargs):\n            command_path = click.get_current_context().command_path\n            try:\n                index = index_connect(local_config,\n                                      application_name=app_name or command_path,\n                                      validate_connection=expect_initialised)\n                _LOG.debug(\"Connected to datacube index: %s\", index)\n            except (OperationalError, ProgrammingError) as e:\n                handle_exception('Error Connecting to database: %s', e)\n                return\n\n            try:\n                return f(index, *args, **kwargs)\n            finally:\n                index.close()\n                del index\n\n        return functools.update_wrapper(with_index, f)\n\n    return decorate\n\n\ndef pass_datacube(app_name=None, expect_initialised=True):\n    \"\"\"\n    Get a DataCube from the current or default local settings.\n\n    :param str app_name:\n        A short name of the application for logging purposes.\n    :param bool expect_initialised:\n        Whether to connect immediately on startup. Useful to catch connection config issues immediately,\n        but if you're planning to fork before any usage (such as in the case of some web servers),\n        you may not want this. For More information on thread/process usage, see datacube.index.Index\n    \"\"\"\n\n    def decorate(f):\n        @pass_index(app_name=app_name, expect_initialised=expect_initialised)\n        def with_datacube(index, *args, **kwargs):\n            return f(Datacube(index=index), *args, **kwargs)\n\n        return functools.update_wrapper(with_datacube, f)\n\n    return decorate\n\n\ndef parse_endpoint(value):\n    ip, port = tuple(value.split(':'))\n    return ip, int(port)\n\n\nEXECUTOR_TYPES = {\n    'serial': lambda _: get_executor(None, None),\n    'multiproc': lambda workers: get_executor(None, int(workers)),\n    'distributed': lambda addr: get_executor(parse_endpoint(addr), True),\n    'celery': lambda addr: mk_celery_executor(*parse_endpoint(addr))\n}\n\nEXECUTOR_TYPES['dask'] = EXECUTOR_TYPES['distributed']  # Add alias \"dask\" for distributed\n\n\ndef _setup_executor(ctx, param, value):\n    try:\n        return EXECUTOR_TYPES[value[0]](value[1])\n    except ValueError:\n        ctx.fail(\"Failed to create '%s' executor with '%s'\" % value)\n\n\nexecutor_cli_options = click.option('--executor',  # type: ignore\n                                    type=(click.Choice(list(EXECUTOR_TYPES)), str),\n                                    default=('serial', None),\n                                    help=\"Run parallelized, either locally or distributed. eg:\\n\"\n                                         \"--executor multiproc 4 (OR)\\n\"\n                                         \"--executor distributed 10.0.0.8:8888\",\n                                    callback=_setup_executor)\n\n\ndef handle_exception(msg, e):\n    \"\"\"\n    Exit following an exception in a CLI app\n\n    If verbosity (-v flag) specified, dump out a stack trace. Otherwise,\n    simply print the given error message.\n\n    Include a '%s' in the message to print the single line message from the\n    exception.\n\n    :param e: caught Exception\n    :param msg: Message to User with optional %s\n    \"\"\"\n    ctx = click.get_current_context()\n    if ctx.obj['verbosity'] >= 1:\n        raise e\n    else:\n        if '%s' in msg:\n            click.echo(msg % e)\n        else:\n            click.echo(msg)\n        ctx.exit(1)\n\n\ndef parsed_search_expressions(f):\n    \"\"\"\n    Add [EXPRESSIONs] arguments to a click application\n\n    Passes a parsed dict of search expressions to the `expressions` argument\n    of the command.\n\n    Also appends documentation on using search expressions to the command.\n\n    WARNING: This wrapped expects an unlimited number of search expressions\n    as click arguments, which means your command must take only click options\n    or a specified number of arguments.\n    \"\"\"\n    if not f.__doc__:\n        f.__doc__ = \"\"\n    f.__doc__ += \"\"\"\n    EXPRESSIONS\n    \n    Select datasets using [EXPRESSIONS] to filter by date, product type,\n    spatial extents or other searchable fields.\n\n    \\b\n        FIELD = VALUE\n        FIELD in DATE-RANGE\n        FIELD in [START, END]\n\n    \\b\n    DATE-RANGE is one of YYYY, YYYY-MM or YYYY-MM-DD\n    START and END can be either numbers or dates\n\n    FIELD: x, y, lat, lon, time, product, ... \n\n    \\b\n    eg. 'time in [1996-01-01, 1996-12-31]'\n        'time in 1996'\n        'lon in [130, 140]' 'lat in [-40, -30]'\n        product=ls5_nbar_albers\n\n    \"\"\"\n\n    def my_parse(ctx, param, value):\n        return parse_expressions(*list(value))\n\n    f = click.argument('expressions', callback=my_parse, nargs=-1)(f)\n    return f\n", "path": "datacube/ui/click.py"}]}