Skip to content

Commit

Permalink
load mmdb from official URL with accountId + licenseKey
Browse files Browse the repository at this point in the history
  • Loading branch information
EvgeniiMunin committed Dec 18, 2024
1 parent 90d3a2e commit 4f9483a
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 30 deletions.
Original file line number Diff line number Diff line change
@@ -1,52 +1,48 @@
package org.prebid.server.hooks.modules.greenbids.real.time.data.config;

import com.google.cloud.storage.Blob;
import com.google.cloud.storage.Storage;
import com.google.cloud.storage.StorageException;
import com.maxmind.db.Reader;
import io.vertx.core.Promise;
import io.vertx.core.Vertx;
import com.maxmind.geoip2.DatabaseReader;
import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
import org.prebid.server.exception.PreBidException;
import org.prebid.server.vertx.Initializable;

import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Optional;
import java.nio.file.StandardCopyOption;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicReference;
import java.util.zip.GZIPInputStream;

public class DatabaseReaderFactory implements Initializable {

private final String gcsBucketName;

private final String geoLiteCountryPath;
private final GreenbidsRealTimeDataProperties properties;

private final Vertx vertx;

private final Storage storage;

private final AtomicReference<DatabaseReader> databaseReaderRef = new AtomicReference<>();

public DatabaseReaderFactory(String gcsBucketName, String geoLiteCountryPath, Vertx vertx, Storage storage) {
this.gcsBucketName = gcsBucketName;
this.geoLiteCountryPath = geoLiteCountryPath;
public DatabaseReaderFactory(GreenbidsRealTimeDataProperties properties, Vertx vertx, Storage storage) {
this.properties = properties;
this.vertx = vertx;
this.storage = storage;
}

@Override
public void initialize(Promise<Void> initializePromise) {
vertx.executeBlocking(() -> {
try {
final Blob blob = getBlob();
final Path databasePath = Files.createTempFile("GeoLite2-Country", ".mmdb");

try (FileOutputStream outputStream = new FileOutputStream(databasePath.toFile())) {
outputStream.write(blob.getContent());
}

databaseReaderRef.set(new DatabaseReader.Builder(databasePath.toFile()).build());
final Path tarGzPath = downloadFile(
properties.geoLiteCountryPath,
properties.maxMindAccountId,
properties.maxMindLicenseKey);
databaseReaderRef.set(extractMMDB(tarGzPath));
} catch (IOException e) {
throw new PreBidException("Failed to initialize DatabaseReader from URL", e);
}
Expand All @@ -55,13 +51,38 @@ public void initialize(Promise<Void> initializePromise) {
.onComplete(initializePromise);
}

private Blob getBlob() {
try {
return Optional.ofNullable(storage.get(gcsBucketName))
.map(bucket -> bucket.get(geoLiteCountryPath))
.orElseThrow(() -> new PreBidException("Bucket not found: " + gcsBucketName));
} catch (StorageException e) {
throw new PreBidException("Error accessing GCS artefact for model: ", e);
private Path downloadFile(String url, String accountId, String licenseKey) throws IOException {
final URL downloadUrl = new URL(url + "&account_id=" + accountId + "&license_key=" + licenseKey);
final HttpURLConnection connection = (HttpURLConnection) downloadUrl.openConnection();
connection.setRequestMethod("GET");

final String auth = accountId + ":" + licenseKey;
final String encodedAuth = Base64.getEncoder().encodeToString(auth.getBytes());
connection.setRequestProperty("Authorization", "Basic " + encodedAuth);

final Path tempFile = Files.createTempFile("geolite2", ".tar.gz");
try (InputStream inputStream = connection.getInputStream()) {
Files.copy(inputStream, tempFile, StandardCopyOption.REPLACE_EXISTING);
}
return tempFile;
}

private DatabaseReader extractMMDB(Path tarGzPath) throws IOException {
try (GZIPInputStream gis = new GZIPInputStream(Files.newInputStream(tarGzPath));
TarArchiveInputStream tarInput = new TarArchiveInputStream(gis)) {

TarArchiveEntry currentEntry;
while ((currentEntry = tarInput.getNextTarEntry()) != null) {
if (currentEntry.getName().contains("GeoLite2-Country.mmdb")) {
break;
}
}

final DatabaseReader databaseReader = new DatabaseReader.Builder(tarInput)
.fileMode(Reader.FileMode.MEMORY).build();
return databaseReader;
} catch (IOException e) {
throw new RuntimeException("Failed to extract MMDB file", e);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,7 @@ public class GreenbidsRealTimeDataConfiguration {
@Bean
DatabaseReaderFactory databaseReaderFactory(
GreenbidsRealTimeDataProperties properties, Vertx vertx, Storage storage) {
return new DatabaseReaderFactory(
properties.getGcsBucketName(), properties.getGeoLiteCountryPath(), vertx, storage);
return new DatabaseReaderFactory(properties, vertx, storage);
}

@Bean
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,10 @@ public class GreenbidsRealTimeDataProperties {

String geoLiteCountryPath;

String maxMindAccountId;

String maxMindLicenseKey;

String gcsBucketName;

Integer cacheExpirationMinutes;
Expand Down

0 comments on commit 4f9483a

Please sign in to comment.