text: #!/usr/bin/env python
  # Copyright 2016 Criteo
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """A CLI to import whisper data into Cassandra."""
  
  from __future__ import print_function
  
  import argparse
  import datetime
  import io
  import logging
  import multiprocessing
  import os
  import re
  import struct
  import sys
  import time
  from multiprocessing import dummy as multiprocessing_dummy
  
  import progressbar
  import scandir
  import whisper
  
  from biggraphite import accessor_factory as bg_accessor_factory
  from biggraphite import metric as bg_metric
  from biggraphite import settings as bg_settings
  from biggraphite import utils as bg_utils
  from biggraphite.cli import command
  
  _DEV_NULL = open(os.devnull, "w")
  
  
  _POINT_STRUCT = struct.Struct(whisper.pointFormat)
  _WORKER = None
  
  
  log = logging.getLogger(__name__)
  
  
  def metric_name_from_wsp(root_dir, prefix, wsp_path):
      """Return the name of a metric given a wsp file path and a root directory.
  
      The path do not have to exist.
  
      Args:
        root_dir: A directory that is parent to all metrics.
        prefix: Prefix to preprend to metric names.
        wsp_path: The name of a file ending with wsp file in root_dir.
  
      Returns:
        The metric name.
      """
      relpath = os.path.relpath(wsp_path, root_dir)
      assert ".." not in relpath, "%s not a child of %" % (root_dir, wsp_path)
      relpath_noext = os.path.splitext(relpath)[0]
      return prefix + relpath_noext.replace(os.path.sep, ".")
  
  
  class _Worker(object):
      def __init__(self, opts):
          settings = bg_settings.settings_from_args(opts)
          bg_utils.set_log_level(settings)
          self._accessor = bg_accessor_factory.accessor_from_settings(settings)
          self._opts = opts
          self.time_start = time.mktime(self._opts.time_start.timetuple())
          self.time_end = time.mktime(self._opts.time_end.timetuple())
  
      @staticmethod
      def _read_metadata(metric_name, path):
          info = whisper.info(path)
          if not info:
              return None
  
          retentions = bg_metric.Retention(
              [
                  bg_metric.Stage(precision=a["secondsPerPoint"], points=a["points"])
                  for a in info["archives"]
              ]
          )
          aggregator = bg_metric.Aggregator.from_carbon_name(info["aggregationMethod"])
          return bg_metric.MetricMetadata.create(
              aggregator=aggregator,
              retention=retentions,
              carbon_xfilesfactor=info["xFilesFactor"],
          )
  
      def _read_points(self, path):
          """Return a list of (timestamp, value)."""
          info = whisper.info(path)
          res = []
          if not info:
              return []
  
          archives = info["archives"]
          with io.open(path, "rb") as f:
              buf = f.read()
  
          stage0 = True
          for archive in archives:
              offset = archive["offset"]
              stage = bg_metric.Stage(
                  precision=archive["secondsPerPoint"],
                  points=archive["points"],
                  stage0=stage0,
              )
              stage0 = False
              if stage in self._opts.ignored_stages:
                  continue
  
              for _ in range(archive["points"]):
                  timestamp, value = _POINT_STRUCT.unpack_from(buf, offset)
                  offset += whisper.pointSize
  
                  if timestamp == 0:
                      continue
                  elif timestamp >= self.time_start and timestamp <= self.time_end:
                      res.append((timestamp, value, 1, stage))
  
          return res
  
      def import_whisper(self, path):
          if not self._accessor.is_connected:
              self._accessor.connect()
  
          name = metric_name_from_wsp(self._opts.root_directory, self._opts.prefix, path)
          metadata = self._read_metadata(name, path)
          log.debug("%s: %s" % (name, metadata.as_string_dict()))
  
          if not metadata:
              return 0
  
          metric = bg_metric.make_metric(name, metadata)
          if not self._opts.no_metadata:
              self._accessor.create_metric(metric)
  
          ret = 0
          if not self._opts.no_data:
              points = self._read_points(path)
              self._accessor.insert_downsampled_points(metric, points)
              ret = len(points)
  
          return ret
  
  
  def _setup_process(opts):
      global _WORKER
      _WORKER = _Worker(opts)
  
  
  def _import_whisper(*args, **kwargs):
      assert _WORKER is not None, "_setup_process was never called"
      try:
          return _WORKER.import_whisper(*args, **kwargs)
      except Exception as e:
          log.exception(e)
          return 0
  
  
  def _parse_opts(args):
      parser = argparse.ArgumentParser(
          description="Import whisper files into BigGraphite."
      )
      parser.add_argument(
          "root_directory",
          metavar="WHISPER_DIR",
          help="directory in which to find whisper files",
      )
      parser.add_argument(
          "--filter",
          type=str,
          default=r".*\.wsp",
          help="Only import metrics matching this filter",
      )
      parser.add_argument(
          "--prefix",
          metavar="WHISPER_PREFIX",
          default="",
          help="prefix to prepend to metric names",
      )
      parser.add_argument(
          "--quiet",
          action="store_const",
          default=False,
          const=True,
          help="Show no output unless there are problems.",
      )
      parser.add_argument(
          "--process",
          metavar="N",
          type=int,
          help="number of concurrent process",
          default=multiprocessing.cpu_count(),
      )
      parser.add_argument(
          "--no-data", action="store_true", help="Do not import data, only metadata."
      )
      parser.add_argument(
          "--no-metadata", action="store_true", help="Do not import metadata, only data."
      )
      parser.add_argument(
          "--ignored_stages",
          nargs="*",
          help="Do not import data for these stages.",
          default=[],
      )
      parser.add_argument(
          "--time-start",
          action=command.ParseDateTimeArg,
          help="Read points written later than this time.",
          default=datetime.datetime.fromtimestamp(0),
          required=False,
      )
      parser.add_argument(
          "--time-end",
          action=command.ParseDateTimeArg,
          help="Read points written earlier than this time.",
          default=datetime.datetime.now(),
          required=False,
      )
      bg_settings.add_argparse_arguments(parser)
      opts = parser.parse_args(args)
      opts.ignored_stages = [bg_metric.Stage.from_string(s) for s in opts.ignored_stages]
      return opts
  
  
  # TODO: put that in a thread.
  class _Walker:
      def __init__(self, root_directory, regexp):
          self.count = 0
          self.root_directory = root_directory
          self.regexp = re.compile(regexp)
  
      def paths(self, root=None):
          root = root or self.root_directory
          for entry in scandir.scandir(root):
              if entry.is_dir():
                  for filename in self.paths(entry.path):
                      yield filename
              elif self.regexp.match(os.path.join(root, entry.name)):
                  self.count += 1
                  yield os.path.join(root, entry.name)
  
  
  def main(args=None):
      """Entry point for the module."""
      if not args:
          args = sys.argv[1:]
  
      opts = _parse_opts(args)
  
      pool_factory = multiprocessing.Pool
      if opts.process == 1:
          pool_factory = multiprocessing_dummy.Pool
      pool = pool_factory(opts.process, initializer=_setup_process, initargs=(opts,))
  
      out_fd = sys.stderr
      if opts.quiet:
          out_fd = _DEV_NULL
      if "__pypy__" not in sys.builtin_module_names:
          print("Running without PyPy, this is about 20 times slower", file=out_fd)
          out_fd.flush()
  
      walker = _Walker(opts.root_directory, opts.filter)
      paths = walker.paths()
      total_points = 0
      max_value = progressbar.UnknownLength
      with progressbar.ProgressBar(
          max_value=max_value, fd=out_fd, redirect_stderr=True
      ) as pbar:
          try:
              res = pool.imap_unordered(_import_whisper, paths)
              for n_path, n_points in enumerate(res):
                  total_points += n_points
                  pbar.update(n_path)
          except KeyboardInterrupt:
              pool.terminate()
  
      pool.close()
      pool.join()
  
      print(
          "Uploaded",
          walker.count,
          "metrics containing",
          total_points,
          "points",
          file=out_fd,
      )
  
  
  if __name__ == "__main__":
      main()
  