WIP

raw-labs · Aug 3, 2024 · 633a3f8 · 633a3f8
1 parent 6688921
commit 633a3f8
Show file tree

Hide file tree

Showing 10 changed files with 86 additions and 49 deletions.
diff --git a/snapi-truffle/src/main/java/raw/runtime/truffle/RawContext.java b/snapi-truffle/src/main/java/raw/runtime/truffle/RawContext.java
@@ -21,9 +21,11 @@
 import java.util.Map;
 import raw.client.api.*;
 import raw.inferrer.api.InferrerService;
+import raw.runtime.truffle.runtime.exceptions.RawTruffleRuntimeException;
 import raw.runtime.truffle.runtime.function.RawFunctionRegistry;
 import raw.utils.AuthenticatedUser;
 import raw.utils.RawSettings;
+import scala.Option;
 import scala.collection.JavaConverters;
 
 public final class RawContext {
@@ -97,12 +99,47 @@ public RawSettings getSettings() {
 
   @CompilerDirectives.TruffleBoundary
   public Map<String, String> getHttpHeaders(String name) {
-    return JavaConverters.mapAsJavaMap(programEnvironment.httpHeaders().get(name).get());
+    Option<scala.collection.immutable.Map<String, String>> maybeHttpHeaders = programEnvironment.httpHeaders().get(name);
+    if (maybeHttpHeaders.isEmpty()) {
+      throw new RawTruffleRuntimeException("unknown http credential: " + name);
+    }
+    return JavaConverters.mapAsJavaMap(maybeHttpHeaders.get());
   }
 
   @CompilerDirectives.TruffleBoundary
   public String getSecret(String key) {
-    return programEnvironment.secrets().get(key).get();
+    Option<String> maybeSecret = programEnvironment.secrets().get(key);
+    if (maybeSecret.isEmpty()) {
+      throw new RawTruffleRuntimeException("unknown secret: " + key);
+    }
+    return maybeSecret.get();
+  }
+
+  @CompilerDirectives.TruffleBoundary
+  public JdbcLocation getJdbcLocation(String name) {
+    Option<JdbcLocation> maybeJdbcLocation = programEnvironment.jdbcServers().get(name);
+    if (maybeJdbcLocation.isEmpty()) {
+      throw new RawTruffleRuntimeException("unknown database credential: " + name);
+    }
+    return maybeJdbcLocation.get();
+  }
+
+  @CompilerDirectives.TruffleBoundary
+  public boolean existsS3Credential(String bucket) {
+    if (programEnvironment.s3Credentials().contains(bucket)) {
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  @CompilerDirectives.TruffleBoundary
+  public S3Credential getS3Credential(String bucket) {
+    Option<S3Credential> maybeCred = programEnvironment.s3Credentials().get(bucket);
+    if (maybeCred.isEmpty()) {
+      throw new RawTruffleRuntimeException("unknown S3 bucket: " + bucket);
+    }
+    return maybeCred.get();
   }
 
   @CompilerDirectives.TruffleBoundary

diff --git a/...ime/truffle/ast/expressions/builtin/location_package/LocationFromMySQLCredentialNode.java b/...ime/truffle/ast/expressions/builtin/location_package/LocationFromMySQLCredentialNode.java
@@ -37,7 +37,7 @@ public Object executeGeneric(VirtualFrame frame) {
     RawContext context = RawContext.get(this);
 
     String credentialName = (String) this.credentialName.executeGeneric(frame);
-    Location l = (Location) context.getProgramEnvironment().jdbcServers().get(credentialName).get();
+    Location l = (Location) context.getJdbcLocation(credentialName);
     MySqlServerLocationDescription d =
         (MySqlServerLocationDescription) LocationDescription$.MODULE$.toLocationDescription(l);
 

diff --git a/...me/truffle/ast/expressions/builtin/location_package/LocationFromOracleCredentialNode.java b/...me/truffle/ast/expressions/builtin/location_package/LocationFromOracleCredentialNode.java
@@ -37,7 +37,7 @@ public Object executeGeneric(VirtualFrame frame) {
     RawContext context = RawContext.get(this);
 
     String credentialName = (String) this.credentialName.executeGeneric(frame);
-    Location l = (Location) context.getProgramEnvironment().jdbcServers().get(credentialName).get();
+    Location l = (Location) context.getJdbcLocation(credentialName);
     OracleServerLocationDescription d =
         (OracleServerLocationDescription) LocationDescription$.MODULE$.toLocationDescription(l);
 

diff --git a/...ruffle/ast/expressions/builtin/location_package/LocationFromPostgreSQLCredentialNode.java b/...ruffle/ast/expressions/builtin/location_package/LocationFromPostgreSQLCredentialNode.java
@@ -37,7 +37,7 @@ public Object executeGeneric(VirtualFrame frame) {
     RawContext context = RawContext.get(this);
 
     String credentialName = (String) this.credentialName.executeGeneric(frame);
-    Location l = (Location) context.getProgramEnvironment().jdbcServers().get(credentialName).get();
+    Location l = (Location) context.getJdbcLocation(credentialName);
     PostgresqlServerLocationDescription d =
         (PostgresqlServerLocationDescription) LocationDescription$.MODULE$.toLocationDescription(l);
 

diff --git a/...java/raw/runtime/truffle/ast/expressions/builtin/location_package/LocationFromS3Node.java b/...java/raw/runtime/truffle/ast/expressions/builtin/location_package/LocationFromS3Node.java
@@ -67,19 +67,19 @@ public Object executeGeneric(VirtualFrame frame) {
     // The docs say:
     // "If the S3 bucket is not registered in the credentials storage, then the region, accessKey
     // and secretKey must be provided as arguments."
-    Option<S3Credential> maybeCred =
-        RawContext.get(this).getProgramEnvironment().s3Credentials().get(bucket);
+    RawContext context = RawContext.get(this);
     S3Path location;
-    if (maybeCred.isDefined()) {
-      S3Credential cred = maybeCred.get();
+
+    if (RawContext.get(this).existsS3Credential(bucket)) {
+      S3Credential cred = context.getS3Credential(bucket);
       location =
           new S3Path(
               bucket,
               cred.region(),
               cred.accessKey(),
               cred.secretKey(),
               path,
-              RawContext.get(this).getSettings());
+              context.getSettings());
     } else {
       // We actually do NOT throw an exception if the accessKey is not passed.
       // Instead, we go without it, which triggers anonymous access to the S3 bucket.
@@ -96,7 +96,7 @@ public Object executeGeneric(VirtualFrame frame) {
               maybeAccessKey,
               maybeSecretKey,
               path,
-              RawContext.get(this).getSettings());
+              context.getSettings());
     }
 
     return new LocationObject(location);

diff --git a/...truffle/ast/expressions/builtin/location_package/LocationFromSQLServerCredentialNode.java b/...truffle/ast/expressions/builtin/location_package/LocationFromSQLServerCredentialNode.java
@@ -37,7 +37,7 @@ public Object executeGeneric(VirtualFrame frame) {
     RawContext context = RawContext.get(this);
 
     String credentialName = (String) this.credentialName.executeGeneric(frame);
-    Location l = (Location) context.getProgramEnvironment().jdbcServers().get(credentialName).get();
+    Location l = (Location) context.getJdbcLocation(credentialName);
     SqlServerServerLocationDescription d =
         (SqlServerServerLocationDescription) LocationDescription$.MODULE$.toLocationDescription(l);
 

diff --git a/...truffle/ast/expressions/builtin/location_package/LocationFromSnowflakeCredentialNode.java b/...truffle/ast/expressions/builtin/location_package/LocationFromSnowflakeCredentialNode.java
@@ -37,7 +37,7 @@ public Object executeGeneric(VirtualFrame frame) {
     RawContext context = RawContext.get(this);
 
     String credentialName = (String) this.credentialName.executeGeneric(frame);
-    Location l = (Location) context.getProgramEnvironment().jdbcServers().get(credentialName).get();
+    Location l = (Location) context.getJdbcLocation(credentialName);
     SnowflakeServerLocationDescription d =
         (SnowflakeServerLocationDescription) LocationDescription$.MODULE$.toLocationDescription(l);
 

diff --git a/.../raw/runtime/truffle/ast/expressions/builtin/location_package/LocationFromStringNode.java b/.../raw/runtime/truffle/ast/expressions/builtin/location_package/LocationFromStringNode.java
@@ -53,14 +53,15 @@ public Object executeGeneric(VirtualFrame frame) {
       throw new RawTruffleRuntimeException("missing protocol: " + url);
     }
 
-    RawSettings rawSettings = RawContext.get(this).getSettings();
+    RawContext context = RawContext.get(this);
+    RawSettings rawSettings = context.getSettings();
 
     Location location = null;
 
     String protocol = url.substring(0, colonIndex);
     switch (protocol) {
       case "pgsql":
-        location = getPgsqlLocation(url, rawSettings);
+        location = getPgsqlLocation(url, context);
         break;
       case "mysql":
         throw new RawTruffleRuntimeException("mysql location not supported");
@@ -81,7 +82,7 @@ public Object executeGeneric(VirtualFrame frame) {
         throw new RawTruffleRuntimeException("snowflake location not supported");
         //        break;
       case "s3":
-        location = getS3Location(url, rawSettings);
+        location = getS3Location(url, context);
         break;
       case "http":
         throw new RawTruffleRuntimeException("http location not supported");
@@ -120,7 +121,7 @@ private Location getLocalLocation(String url) {
   }
 
   @CompilerDirectives.TruffleBoundary
-  private Location getS3Location(String url, RawSettings settings) {
+  private Location getS3Location(String url, RawContext context) {
     try {
       URI uri = new URI(url);
       String uriUserInfo = uri.getUserInfo();
@@ -130,7 +131,6 @@ private Location getS3Location(String url, RawSettings settings) {
 
       String accessKey = null;
       String secretKey = null;
-      String region = null;
       if (uriUserInfo != null) {
         String[] userInfoParts = uriUserInfo.split(":");
         accessKey = userInfoParts[0];
@@ -150,15 +150,14 @@ private Location getS3Location(String url, RawSettings settings) {
       if (accessKey == null) {
         // If the access key/secret key are not defined, then the "host" is actually the bucket name
         // in the program environment credentials set.
-        S3Credential s3Credential =
-            RawContext.get(this).getProgramEnvironment().s3Credentials().get(bucketName).get();
+        S3Credential s3Credential = RawContext.get(this).getS3Credential(bucketName);
         return new S3Path(
             bucketName,
             s3Credential.region(),
             s3Credential.accessKey(),
             s3Credential.secretKey(),
             objectKey,
-            settings);
+            context.getSettings());
       } else {
         // TODO (msb): There is no way to specify the region when using a direct URL...
         return new S3Path(
@@ -167,15 +166,15 @@ private Location getS3Location(String url, RawSettings settings) {
             new Some(accessKey),
             new Some(secretKey),
             objectKey,
-            settings);
+            context.getSettings());
       }
     } catch (URISyntaxException e) {
       throw new RawTruffleRuntimeException("invalid S3 URL: " + url);
     }
   }
 
   @CompilerDirectives.TruffleBoundary
-  private Location getPgsqlLocation(String url, RawSettings settings) {
+  private Location getPgsqlLocation(String url, RawContext context) {
     try {
       URI uri = new URI(url);
       String uriUserInfo = uri.getUserInfo();
@@ -237,12 +236,12 @@ private Location getPgsqlLocation(String url, RawSettings settings) {
 
       if (dbname != null && schema != null && table != null) {
         return new PostgresqlTableLocation(
-            host, port, dbname, username, password, schema, table, settings);
+            host, port, dbname, username, password, schema, table, context.getSettings());
       } else if (dbname != null && schema != null) {
         return new PostgresqlSchemaLocation(
-            host, port, dbname, username, password, schema, settings);
+            host, port, dbname, username, password, schema, context.getSettings());
       } else if (dbname != null) {
-        return new PostgresqlServerLocation(host, port, dbname, username, password, settings);
+        return new PostgresqlServerLocation(host, port, dbname, username, password, context.getSettings());
       } else {
         throw new RawTruffleRuntimeException("invalid PostgreSQL URL: " + url);
       }

diff --git a/...i-truffle/src/main/java/raw/runtime/truffle/ast/expressions/literals/BinaryConstNode.java b/...i-truffle/src/main/java/raw/runtime/truffle/ast/expressions/literals/BinaryConstNode.java
@@ -18,19 +18,19 @@
 
 public class BinaryConstNode extends ExpressionNode {
 
-    private final BinaryObject value;
+  private final BinaryObject value;
 
-    public BinaryConstNode(byte[] value) {
-        this.value = new BinaryObject(value);
-    }
+  public BinaryConstNode(byte[] value) {
+    this.value = new BinaryObject(value);
+  }
 
-    @Override
-    public final BinaryObject executeBinary(VirtualFrame virtualFrame) {
-        return value;
-    }
+  @Override
+  public final BinaryObject executeBinary(VirtualFrame virtualFrame) {
+    return value;
+  }
 
-    @Override
-    public final Object executeGeneric(VirtualFrame virtualFrame) {
-        return value;
-    }
+  @Override
+  public final Object executeGeneric(VirtualFrame virtualFrame) {
+    return value;
+  }
 }
diff --git a/...truffle/src/main/java/raw/runtime/truffle/ast/expressions/literals/LocationConstNode.java b/...truffle/src/main/java/raw/runtime/truffle/ast/expressions/literals/LocationConstNode.java
@@ -16,23 +16,24 @@
 import raw.compiler.rql2.api.LocationDescription$;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawContext;
-import raw.runtime.truffle.runtime.primitives.BinaryObject;
 import raw.runtime.truffle.runtime.primitives.LocationObject;
 import raw.sources.api.Location;
 import raw.utils.RawSettings;
 
 public class LocationConstNode extends ExpressionNode {
 
-    private final byte[] value;
+  private final byte[] value;
 
-    public LocationConstNode(byte[] value) {
-        this.value = value;
-    }
+  public LocationConstNode(byte[] value) {
+    this.value = value;
+  }
 
-    @Override
-    public final Object executeGeneric(VirtualFrame virtualFrame) {
-        RawSettings rawSettings = RawContext.get(this).getSettings();
-        Location location = LocationDescription$.MODULE$.toLocation(LocationDescription$.MODULE$.deserialize(value), rawSettings);
-        return new LocationObject(location);
-    }
+  @Override
+  public final Object executeGeneric(VirtualFrame virtualFrame) {
+    RawSettings rawSettings = RawContext.get(this).getSettings();
+    Location location =
+        LocationDescription$.MODULE$.toLocation(
+            LocationDescription$.MODULE$.deserialize(value), rawSettings);
+    return new LocationObject(location);
+  }
 }