def test_export_to_file_append(self):
exporter = MagicMock()
product = {
'exporter': exporter,
'format': 'the format',
'append': True,
'append_to_filename': 'the filename',
}
# Do append, pass value of append_to_filename to exporter
self.assertEqual(
exporter(),
export_to_file('host', product, 'file', 'catalogue', 'collection'))
exporter.assert_called_with(ANY,
'file',
'the format',
append='the filename',
filter=None)
# Don't append
product['append'] = False
self.assertEqual(
exporter(),
export_to_file('host', product, 'file', 'catalogue', 'collection'))
exporter.assert_called_with(ANY,
'file',
'the format',
append=False,
filter=None)
def test_export_encrypt_file(self, mock_encrypt_file):
exporter = MagicMock()
product = {
'exporter': exporter,
'format': 'the format',
'encryption_key': 'any key'
}
# Do append, pass value of append_to_filename to exporter
export_to_file('host', product, 'file', 'catalogue', 'collection')
mock_encrypt_file.assert_called_with('file', 'any key')
def test_export_to_file_entity_filters(self, mock_group_filter, mock_buffered_iterable, mock_graphql_streaming):
from gobexport.exporter import export_to_file
mock_entity_filter = MagicMock()
mock_exporter = MagicMock()
product = {
'api_type': 'graphql_streaming',
'query': 'some query',
'exporter': mock_exporter,
'format': 'the format',
'unfold': 'true_or_false',
'entity_filters': [
mock_entity_filter,
]
}
result = export_to_file('host', product, 'file', 'catalogue', 'collection')
mock_filter = mock_group_filter.return_value
mock_group_filter.assert_called_with([mock_entity_filter])
mock_exporter.assert_called_with(mock_buffered_iterable.return_value,
'file',
'the format',
append=False,
filter=mock_group_filter.return_value)
# Assert the reset function is called on the filter
mock_filter.reset.assert_called()
def test_export_to_file(monkeypatch):
global records
monkeypatch.setitem(__builtins__, 'open', mock_open)
monkeypatch.setattr(gobexport.api, 'API', MockAPI)
monkeypatch.setattr(gobexport.graphql, 'GraphQL', MockGraphQL)
before_each(monkeypatch)
from gobexport.exporter import export_to_file
# Test DAT export
catalogue = 'meetbouten'
collection = 'meetbouten'
# Get the configuration for this collection
config = CONFIG_MAPPING[catalogue][collection]
export_to_file('host', config.products['dat'], '/tmp/ttt', catalogue, collection)
assert(MockFile.s == '$$2$$||125,6|10,1||||||||||||||POINT (125.6 10.1)\n')
MockFile.s = ''
catalogue = 'gebieden'
collection = 'stadsdelen'
# Get the configuration for this collection
config = CONFIG_MAPPING[catalogue][collection]
format = config.products['csv_actueel'].get('format')
export_to_file('host', config.products['csv_actueel'], '/tmp/ttt', catalogue, collection)
expected_result = 'identificatie;code;naam;beginGeldigheid;eindGeldigheid;documentdatum;documentnummer;ligtIn:BRK.GME.identificatie;ligtIn:BRK.GME.naam;geometrie\r\n2;;;;;;;;;POINT (125.6 10.1)\r\n'
assert(MockFile.s == expected_result)
MockFile.s = ''
# Get the configuration for this collection
config = MockConfig
export_to_file('host', config.products['csv'], '/tmp/ttt', catalogue, collection)
expected_result = 'identificatie;boolean;geometrie\r\n2;J;POINT (125.6 10.1)\r\n'
assert(MockFile.s == expected_result)
format = config.products['shape'].get('format')
file_name = 'esri.shp'
# Update records to contain an geometry collection
records = [{'identificatie': '2', 'boolean': False, 'geometrie': {'type': 'GeometryCollection', 'geometries': [{'type': 'LineString', 'coordinates': [[125891.16, 480253.38], [125891.07, 480253.34]]}, {'type': 'Polygon', 'coordinates': [[[125891.16, 480253.38], [125893.06, 480250.0], [125892.57, 480250.0]]]}]}}]
export_to_file('host', config.products['shape'], file_name, catalogue, collection)
# Remove created files
for file in ['esri.shp', 'esri.dbf', 'esri.shx', 'esri.prj']:
assert(os.path.isfile(file))
os.remove(file)
def test_export_to_file_objectstore(self, mock_buffered_iterable, mock_objectstore_file):
from gobexport.exporter import export_to_file
product = {
'api_type': 'objectstore',
'exporter': MagicMock(),
'format': 'the format',
'config': 'the config',
}
result = export_to_file('host', product, 'file', 'catalogue', 'collection', False)
mock_objectstore_file.assert_called_with(product['config'], row_formatter=None)
mock_buffered_iterable.assert_called_with(mock_objectstore_file.return_value, 'source', buffer_items=False)
product['exporter'].assert_called_with(mock_buffered_iterable.return_value, 'file', 'the format',
append=False, filter=None)
def test_export_to_file_graphql_streaming_with_unfold(self, mock_buffered_iterable, mock_graphql_streaming):
from gobexport.exporter import export_to_file
product = {
'api_type': 'graphql_streaming',
'query': 'some query',
'exporter': MagicMock(),
'format': 'the format',
'unfold': 'true_or_false',
'sort': 'sorter',
'row_formatter': 'row_form',
}
result = export_to_file('host', product, 'file', 'catalogue', 'collection', False)
mock_graphql_streaming.assert_called_with('host', product['query'], unfold='true_or_false', sort='sorter',
row_formatter='row_form', cross_relations=False, batch_size=None,
secure_user=None)
def test_export_to_file_graphql_streaming(self, mock_buffered_iterable, mock_graphql_streaming):
from gobexport.exporter import export_to_file
product = {
'api_type': 'graphql_streaming',
'secure_user': '******',
'query': 'some query',
'exporter': MagicMock(),
'format': 'the format',
}
result = export_to_file('host', product, 'file', 'catalogue', 'collection', False)
mock_graphql_streaming.assert_called_with('host', product['query'], row_formatter=None, sort=None,
unfold=False, cross_relations=False, batch_size=None,
secure_user='******')
mock_buffered_iterable.assert_called_with(mock_graphql_streaming.return_value, 'source', buffer_items=False)
product['exporter'].assert_called_with(mock_buffered_iterable.return_value, 'file', 'the format',
append=False, filter=None)
def _export_collection(host, catalogue, collection, product_name,
destination): # noqa: C901
"""Export a collection from a catalog
:param host: The API host to retrieve the catalog and collection from
:param catalog: The name of the catalog
:param collection: The name of the collection
:param product_name: The name of the product to export
:param destination: The destination of the resulting output file(s)
:return:
"""
logger.info(f"Export {catalogue}:{collection} to {destination} started.")
# Get the configuration for this collection
config = CONFIG_MAPPING[catalogue][collection]
resolve_config_filenames(config)
files = []
# If a product has been supplied, export only that product
try:
products = {
product_name: config.products[product_name]
} if product_name else config.products
except KeyError:
logger.error(f"Product '{product_name}' not found")
return
# Start exporting each product
for name, product in products.items():
logger.info(
f"Export to file '{name}' started, API type: {product.get('api_type', 'REST')}"
)
# Get name of local file to write results to
results_file = _get_filename(
product['filename']
) if destination == "Objectstore" else product['filename']
if product.get('append', False):
# Add .to_append to avoid writing to the previously created file
results_file = _get_filename(f"{product['filename']}.to_append")
product['append_to_filename'] = _get_filename(product['filename']) \
if destination == "Objectstore" \
else product['filename']
# Buffer items if they are used multiple times. This prevents calling API multiple times for same data
source = product_source(product)
buffer_items = len(
list(
filter(lambda p: product_source(p) == source,
config.products.values()))) > 1
logger.info(
f"Buffering API output {'enabled' if buffer_items else 'disabled'}"
)
try:
row_count = _with_retries(
lambda: export_to_file(host,
product,
results_file,
catalogue,
product.get('collection', collection),
buffer_items=buffer_items))
except Exception as e:
logger.error(f"Export to local file {name} failed: {str(e)}.")
else:
logger.info(f"{row_count} records exported to local file {name}.")
if product.get('append', False):
# Append temporary file to existing file and cleanup temp file
_append_to_file(results_file, product['append_to_filename'])
os.remove(results_file)
else:
# Do not add file to files again when appending
files.append({
'temp_location': results_file,
'distribution': product['filename'],
'mime_type': product['mime_type']
})
# Add extra result files (e.g. .prj file)
extra_files = product.get('extra_files', [])
files.extend([{
'temp_location': _get_filename(file['filename']),
'distribution': file['filename'],
'mime_type': file['mime_type']
} for file in extra_files])
if destination == "Objectstore":
# Get objectstore connection
config = get_datastore_config(GOB_OBJECTSTORE)
datastore = DatastoreFactory.get_datastore(config)
datastore.connect()
assert isinstance(datastore, ObjectDatastore)
connection = datastore.connection
logger.info(
f"Connection to {destination} {datastore.user} has been made.")
# Start distribution of all resulting files
for file in files:
logger.info(f"Write file '{file['distribution']}'.")
if destination == "Objectstore":
# Distribute to pre-final location
container = f'{CONTAINER_BASE}/{EXPORT_DIR}/{catalogue}/'
with open(file['temp_location'], 'rb') as fp:
try:
distribute_to_objectstore(connection, container,
file['distribution'], fp,
file['mime_type'])
except GOBException as e:
logger.error(
f"Failed to copy to {destination} on location: {container}{file['distribution']}. \
Error: {e}")
return False
logger.info(
f"File copied to {destination} on location: {container}{file['distribution']}."
)
cleanup_datefiles(
connection, CONTAINER_BASE,
f"{EXPORT_DIR}/{catalogue}/{file['distribution']}")
# Delete temp file
os.remove(file['temp_location'])
elif destination == "File":
logger.info(f"Export is written to {file['distribution']}.")
logger.info("Export completed")
