CREATE OR REPLACE NETWORK RULE ext_neo4j_network_rule
  MODE = EGRESS
  TYPE = HOST_PORT
  VALUE_LIST = ('xyz.databases.neo4j.io:7687');


CREATE OR REPLACE SECRET ext_neo4j_credentials
    TYPE = password
    USERNAME = 'xxx'
    PASSWORD = '*******';


CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION ext_neo4j_access_integration
  ALLOWED_NETWORK_RULES = (ext_neo4j_network_rule)
  ALLOWED_AUTHENTICATION_SECRETS=(ext_neo4j_credentials)
  ENABLED = true;


CREATE OR REPLACE FUNCTION read_neo4j(OPTION OBJECT, query STRING)
  RETURNS TABLE(data OBJECT)
  LANGUAGE JAVA
  RUNTIME_VERSION = '17'
  IMPORTS = ('@NEO4J_JARS/drivers/neo4j-jdbc-full-bundle-6.2.0.jar')
  EXTERNAL_ACCESS_INTEGRATIONS = (ext_neo4j_access_integration)
  SECRETS = ('cred' = ext_neo4j_credentials )
  HANDLER = 'JdbcDataReader'
AS $$
import java.io.PrintWriter;
import java.io.StringWriter;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Stream;

import com.snowflake.snowpark_java.types.SnowflakeSecrets;

public class JdbcDataReader {

	static final class OutputRow {
		public final Map<String, String> data;

		OutputRow(Map<String, String> data) {
			this.data = data;
		}
	}

	public static Class<?> getOutputClass() {
		return OutputRow.class;
	}

	public Stream<OutputRow> process(Map<String, String> jdbcConfig, String query) {
		String jdbcUrl = jdbcConfig.get("url");
		String username;
		String password;

		if ("true".equals(jdbcConfig.get("use_secrets"))) {
			SnowflakeSecrets sfSecrets = SnowflakeSecrets.newInstance();
			var secret = sfSecrets.getUsernamePassword("cred");
			username = secret.getUsername();
			password = secret.getPassword();
		} else {
			username = jdbcConfig.get("username");
			password = jdbcConfig.get("password");
		}

		try (var connection = DriverManager.getConnection(jdbcUrl, username, password);
		     var statement = connection.createStatement();
		     var resultSet = statement.executeQuery(query);
		) {
			var metaData = resultSet.getMetaData();
			var columnNames = new ArrayList<String>();
			var columnCount = metaData.getColumnCount();
			for (int i = 1; i <= columnCount; i++) {
				columnNames.add(metaData.getColumnName(i));
			}

			var streamBuilder = Stream.<OutputRow>builder();
			while (resultSet.next()) {
				Map<String, String> rowMap = new HashMap<>();
				for (String columnName : columnNames) {
					String columnValue = resultSet.getString(columnName);
					rowMap.put(columnName, columnValue);
				}
				streamBuilder.add(new OutputRow(rowMap));
			}
			return streamBuilder.build();
		} catch (SQLException ex) {
			ex.printStackTrace();
			var sw = new StringWriter();
			var pw = new PrintWriter(sw);
			ex.printStackTrace(pw);
			pw.flush();
			return Stream.of(new OutputRow(Map.of("ERROR", sw.toString())));
		}
	}
}
$$;


select * from
table(READ_NEO4J(
  {'url': 'jdbc:neo4j+s://xyz.databases.neo4j.io/xyz', 'use_secrets': 'true'},
$$
MATCH (n:Movie)<-[:ACTED_IN]-(p:Person) 
WHERE n.released >= 2000
RETURN p.name, n.title AS movies
ORDER BY p.name
$$));


select * from
table(READ_NEO4J(
  {'url': 'jdbc:neo4j+s://xyz.databases.neo4j.io/databasename?enableSQLTranslation=true', 'use_secrets': 'true'},
$$
select p.name as name, m.title as title
from Person p natural join ACTED_IN natural join Movie m
order by p.name
$$));