def test_manager_context_manager(monkeypatch):
class MockConnection(object):
def __init__(self):
self.close_called = False
def close(self):
self.close_called = True
def mockreturn(*args, **kwargs):
return MockConnection()
monkeypatch.setattr(sqlalchemy.engine.Engine, 'connect', mockreturn)
factory = database.ManagerFactory('postgres', 'table_name')
postgres_standard = factory.get_manager(use_raster=False)
with postgres_standard:
assert postgres_standard.connection.close_called is False
assert isinstance(postgres_standard.connection, MockConnection)
assert postgres_standard.connection.close_called is True
postgres_raster = factory.get_manager(use_raster=True)
with postgres_raster:
assert postgres_raster.connection.close_called is False
assert isinstance(postgres_raster.connection, MockConnection)
assert postgres_raster.connection.close_called is True
def test_manager_factory_get_manager(monkeypatch):
factory = database.ManagerFactory('postgres', 'table_name')
postgres_standard = factory.get_manager(use_raster=False)
assert isinstance(postgres_standard, database.PostgresValueManager)
assert isinstance(postgres_standard.engine, sqlalchemy.engine.Engine)
assert postgres_standard.connection is None
assert postgres_standard.table_name == 'table_name'
postgres_raster = factory.get_manager(use_raster=True)
assert isinstance(postgres_raster, database.PostgresRasterManager)
assert isinstance(postgres_raster.engine, sqlalchemy.engine.Engine)
assert postgres_raster.connection is None
assert postgres_raster.table_name == 'table_name'
with pytest.raises(Exception) as e:
factory = database.ManagerFactory('couchdb', 'table_name')
factory.get_manager()
assert str(e.value) == "Can't load plugin: sqlalchemy.dialects:couchdb"
def test_manager_factory_constructor_call_private_create_engine_method(
monkeypatch):
mock_callable = tools.MockCallable()
def mockreturn(*args, **kwargs):
return mock_callable(*args, **kwargs)
monkeypatch.setattr(database.ManagerFactory,
'_ManagerFactory__create_engine', mockreturn)
database.ManagerFactory('postgres',
'table_name',
pool_size=5,
host='localhost',
port=3306)
assert mock_callable.args[1:] == ('postgres', )
assert mock_callable.kwargs == {
'host': 'localhost',
'port': 3306,
'pool_size': 5
}
def test_manager_factory_constructor():
factory = database.ManagerFactory('postgres', 'table_name')
assert isinstance(factory.engine, sqlalchemy.engine.Engine)
assert factory.db_driver == 'postgres'
assert factory.table_name == 'table_name'
def load_hgt():
""" Function called by the console_script `gmalt-hgtload`
Usage:
gmalt-hgtload [options] -u <user> <folder>
"""
# Parse command line arguments
parser = create_load_hgt_parser()
args = vars(parser.parse_args())
# logging
traceback = args.pop('traceback')
tools.configure_logging(args.pop('verbose'), echo=args.pop('echo'))
# Pop everything not related to database uri string
concurrency = args.pop('concurrency')
folder = args.pop('folder')
use_raster = args.pop('use_raster')
samples = args.pop('sample')
db_driver = args.pop('type')
table_name = args.pop('table')
check_raster2pgsql = args.pop('check_raster2pgsql')
# sqlalchemy.engine.url.URL args
db_info = args
# If postgres driver and raster2pgsql is available, propose to use this solution instead.
if db_driver == 'postgres' and use_raster and check_raster2pgsql and tools.check_for_raster2pgsql(
):
sys.exit(0)
logging.info('config - parallelism : %i' % concurrency)
logging.info('config - folder : %s' % folder)
logging.info('config - db driver : %s' % db_driver)
logging.info('config - db host : %s' % db_info.get('host'))
logging.info('config - db user : %s' % db_info.get('username'))
logging.info('config - db name : %s' % db_info.get('database'))
logging.info('config - db table : %s' % table_name)
if use_raster:
logging.debug('config - use raster : %s' % use_raster)
logging.debug('config - raster sampling : {}'.format('{}x{}'.format(
*samples) if samples[0] else 'none'))
# create sqlalchemy engine
factory = database.ManagerFactory(db_driver,
table_name,
pool_size=concurrency,
**db_info)
try:
# First validate that the database is ready
with factory.get_manager(use_raster) as manager:
manager.prepare_environment()
# Then process HGT files
tools.import_hgt_zip_files(folder, concurrency, factory, use_raster,
samples)
except sqlalchemy.exc.OperationalError:
logging.error(
'Unable to connect to database with these settings : {}'.format(
factory.engine.url),
exc_info=traceback)
except database.NotSupportedException:
logging.error(
'Database does not support raster settings. Have you enabled GIS extension ?',
exc_info=traceback)
return sys.exit(1)
except KeyboardInterrupt:
return sys.exit(0)
except worker.WorkerPoolException:
# in case of ThreadPoolException, the worker which raised the error
# logs it using logging.exception
return sys.exit(1)
except Exception as e:
logging.error('Unknown error : {}'.format(str(e)), exc_info=traceback)
return sys.exit(1)
return sys.exit(0)