Add an option to run-crate to choose JAVA_HOME automatically

cr8/java_magic.py

@@ -0,0 +1,66 @@
+import os
+import subprocess
+import re
+from glob import glob
+from pathlib import Path
+from typing import Callable
+
+from cr8.misc import parse_version
+
+
+JAVA_CANDIDATES = tuple(
+    glob('/usr/lib/jvm/java-*-openjdk')
+    + glob('/usr/lib/java-*')
+    + glob('/Library/Java/JavaVirtualMachines/jdk*/Contents/Home')
+)
+
+VERSION_RE = re.compile(r'(\d+\.\d+\.\d+(_\d+)?)|"(\d+)"')
+
+
+def _parse_java_version(line: str) -> tuple:
+    """ Return the version number found in the first line of `java -version`
+
+    >>> _parse_java_version('openjdk version "11.0.2" 2018-10-16')
+    (11, 0, 2)
+    """
+    m = VERSION_RE.search(line)
+    version_str = m and m.group(0).replace('"', '') or '0.0.0'
+    if '_' in version_str:
+        fst, snd = version_str.split('_', maxsplit=2)
+        version = parse_version(fst)
+        return (version[1], version[2], int(snd))
+    else:
+        return parse_version(version_str)
+
+
+def _detect_java_version(java_home: str) -> tuple:
+    p = subprocess.run(
+        [Path(java_home) / 'bin' / 'java', '-version'],
+        check=True,
+        encoding='utf-8',
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT
+    )
+    line = p.stdout.split('\n')[0]
+    assert 'version' in line, ('First line must contain the version, got:' + line)
+    return _parse_java_version(line)
+
+
+def _find_matching_java_home(version_matches: Callable[[], bool]) -> str:
+    java_home = os.environ.get('JAVA_HOME', '')
+    for path in filter(os.path.exists, (java_home, ) + JAVA_CANDIDATES):
+        version = _detect_java_version(path)
+        if version_matches(version):
+            return path
+    return java_home
+
+
+def find_java_home(cratedb_version: tuple) -> str:
+    """ Return a path to a JAVA_HOME suites for the given CrateDB version """
+    if (2, 3) <= cratedb_version < (4, 0):
+        # Supports 8 to 11+, use whatever is set
+        return os.environ.get('JAVA_HOME', '')
+    if cratedb_version < (2, 3):
+        return _find_matching_java_home(lambda ver: ver[0] == 8)
+    else:
+        return _find_matching_java_home(lambda ver: ver[0] >= 11)

cr8/misc.py

@@ -31,15 +31,19 @@ def try_len(o: Any) -> int:
 
 
 def parse_version(version: str) -> tuple:
-    """Parse a string formatted X.Y.Z version number into a tuple
+    """Parse a string formatted X[.Y.Z] version number into a tuple
 
     >>> parse_version('10.2.3')
     (10, 2, 3)
+
+    >>> parse_version('12')
+    (12, 0, 0)
     """
     if not version:
         return None
-    major, minor, patch = version.split('.', maxsplit=3)
-    return (int(major), int(minor), int(patch))
+    parts = version.split('.')
+    missing = 3 - len(parts)
+    return tuple(int(i) for i in parts + ([0] * missing))
 
 
 def parse_table(fq_table: str) -> Tuple[str, str]:

cr8/run_crate.py

@@ -25,6 +25,7 @@
 from typing import Dict, Any, List, NamedTuple
 from urllib.request import urlopen
 
+from cr8.java_magic import find_java_home
 from cr8.misc import parse_version, init_logging
 from cr8.engine import DotDict
 from cr8.exceptions import ArgumentError
@@ -195,21 +196,28 @@ def __init__(self,
                  crate_dir: str,
                  env: Dict[str, Any] = None,
                  settings: Dict[str, Any] = None,
-                 keep_data: bool = False) -> None:
+                 keep_data: bool = False,
+                 java_magic: bool = False) -> None:
         """Create a CrateNode
 
         Args:
             crate_dir: Path to the extracted Crate tarball
             env: Environment variables with which the Crate process will be
                 started.
             settings: Additional Crate settings.
+            java_magic: If set to true, it will attempt to set JAVA_HOME to
+                some path that contains a Java version suited to run the given
+                CrateDB instance.
         """
         super().__init__()
         self.crate_dir = crate_dir
         version = _extract_version(crate_dir)
         self.env = env or {}
-        self.env.setdefault('JAVA_HOME',
-                            os.environ.get('JAVA_HOME', ''))
+        if java_magic:
+            java_home = find_java_home(version)
+        else:
+            java_home = os.environ.get('JAVA_HOME', '')
+        self.env.setdefault('JAVA_HOME', java_home)
         self.env.setdefault('LANG',
                             os.environ.get('LANG', os.environ.get('LC_ALL')))
         if not self.env['LANG']:
@@ -611,7 +619,14 @@ def _parse_options(options: List[str]) -> Dict[str, str]:
             f'Option must be in format <key>=<value>, got: {options}')
 
 
-def create_node(version, env=None, setting=None, crate_root=None, keep_data=False):
+def create_node(
+        version,
+        env=None,
+        setting=None,
+        crate_root=None,
+        keep_data=False,
+        java_magic=False,
+):
     init_logging(log)
     settings = {
         'cluster.name': 'cr8-crate-run' + str(random.randrange(1e9))
@@ -622,7 +637,12 @@ def create_node(version, env=None, setting=None, crate_root=None, keep_data=Fals
     if env:
         env = _parse_options(env)
     return CrateNode(
-        crate_dir=crate_dir, env=env, settings=settings, keep_data=keep_data)
+        crate_dir=crate_dir,
+        env=env,
+        settings=settings,
+        keep_data=keep_data,
+        java_magic=java_magic
+    )
 
 
 @argh.arg('version', help='Crate version to run')
@@ -631,8 +651,18 @@ def create_node(version, env=None, setting=None, crate_root=None, keep_data=Fals
 @argh.arg('-s', '--setting', action='append',
           help='Crate setting. Option can be specified multiple times.')
 @argh.arg('--keep-data', help='If this is set the data folder will be kept.')
+@argh.arg(
+    '--disable-java-magic',
+    help='Disable the logic to detect a suitable JAVA_HOME')
 @argh.wrap_errors([ArgumentError])
-def run_crate(version, env=None, setting=None, crate_root=None, keep_data=False):
+def run_crate(
+        version,
+        env=None,
+        setting=None,
+        crate_root=None,
+        keep_data=False,
+        disable_java_magic=False,
+):
     """Launch a crate instance.
 
     Supported version specifications:
@@ -656,7 +686,14 @@ def run_crate(version, env=None, setting=None, crate_root=None, keep_data=False)
     The postgres host and port are available as {node.addresses.psql.host} and
     {node.addresses.psql.port}
     """
-    with create_node(version, env, setting, crate_root, keep_data) as n:
+    with create_node(
+            version,
+            env,
+            setting,
+            crate_root,
+            keep_data,
+            java_magic=not disable_java_magic,
+    ) as n:
         try:
             n.start()
             n.process.wait()

tests/test_java_magic.py

@@ -0,0 +1,28 @@
+from unittest import TestCase
+from doctest import DocTestSuite
+from cr8 import java_magic
+from cr8.java_magic import _parse_java_version
+
+
+class JavaVersionParsingTest(TestCase):
+
+    def assertVersion(self, line, expected):
+        version = _parse_java_version(line)
+        self.assertEqual(version, expected)
+
+    def test_java_8_line(self):
+        self.assertVersion('openjdk version "1.8.0_202"', (8, 0, 202))
+
+    def test_java_10_line(self):
+        self.assertVersion('openjdk version "10.0.2" 2018-07-17', (10, 0, 2))
+
+    def test_java_11_line(self):
+        self.assertVersion('java 11.0.1 2018-10-16 LTS', (11, 0, 1))
+
+    def test_java_12_line(self):
+        self.assertVersion('openjdk version "12" 2019-03-19', (12, 0, 0))
+
+
+def load_tests(loader, tests, ignore):
+    tests.addTests(DocTestSuite(java_magic))
+    return tests