server/main.py - incubator-ponymail-foal - Git at Google

 #!/usr/bin/env python3
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you 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.

 """Apache Pony Mail, Codename Foal - A Python variant of Pony Mail"""
 import argparse
 import asyncio
 import importlib
 import json
 import os
 import sys
 from time import sleep
 import traceback
 import typing

 import aiohttp.web
 import yaml
 import uuid

 import plugins.background
 import plugins.configuration
 import plugins.database
 import plugins.formdata
 import plugins.offloader
 import plugins.server
 import plugins.session

 PONYMAIL_FOAL_VERSION = "0.1.0"
 from server_version import PONYMAIL_SERVER_VERSION


 # Certain environments such as MinGW-w64 will not register as a TTY and uses buffered output.
 # In such cases, we need to force a flush of each print, or nothing will show.
 if not sys.stdout.buffer.isatty():
     import functools
     print = functools.partial(print, flush=True)


 class Server(plugins.server.BaseServer):
     """Main server class, responsible for handling requests and scheduling offloader threads """

     def _load_endpoint(self, subdir):
         for endpoint_file in sorted(os.listdir(subdir)):
             if endpoint_file.endswith(".py"):
                 endpoint = endpoint_file[:-3]
                 m = importlib.import_module(f"{subdir}.{endpoint}")
                 if hasattr(m, "register"):
                     self.handlers[endpoint] = m.__getattribute__("register")(self)
                     print(f"Registered endpoint /api/{endpoint}")
                 else:
                     print(
                         f"Could not find entry point 'register()' in {endpoint_file}, skipping!"
                     )

     def __init__(self, args: argparse.Namespace):
         print(
             "==== Apache Pony Mail (Foal v/%s ~%s) starting... ====" % (PONYMAIL_FOAL_VERSION, PONYMAIL_SERVER_VERSION)
         )
         # Load configuration
         yml = yaml.safe_load(open(args.config))
         self.config = plugins.configuration.Configuration(yml)
         self.data = plugins.configuration.InterData()
         self.handlers = dict()
         self.dbpool = asyncio.Queue()
         self.runners = plugins.offloader.ExecutorPool()
         self.server = None
         self.streamlock = asyncio.Lock()
         self.api_logger = None
         self.foal_version = PONYMAIL_FOAL_VERSION
         self.server_version = PONYMAIL_SERVER_VERSION
         self.stoppable = False # allow remote stop for tests
         self.background_event = asyncio.Event() # for background task to wait on

         # Make a pool of database connections for async queries
         pool_size = self.config.database.pool_size
         if pool_size < 1:
             raise ValueError(f"pool_size {pool_size} must be > 0")
         for _ in range(0, pool_size): # stop value is exclusive
             self.dbpool.put_nowait(plugins.database.Database(self.config.database))

         # Load each URL endpoint
         if args.testendpoints:
             print("** Loading additional testing endpoints **")
             self._load_endpoint("testendpoints")
             print()
         self._load_endpoint("endpoints")

         if args.logger:
             import logging
             es_logger = logging.getLogger('elasticsearch')
             es_logger.setLevel(args.logger)
             es_logger.addHandler(logging.StreamHandler())
         if args.trace:
             import logging
             es_trace_logger = logging.getLogger('elasticsearch.trace')
             es_trace_logger.setLevel(args.trace)
             es_trace_logger.addHandler(logging.StreamHandler())
         if args.apilog:
             import logging
             self.api_logger = logging.getLogger('ponymail.apilog')
             self.api_logger.setLevel(args.apilog)
             self.api_logger.addHandler(logging.StreamHandler())
         self.stoppable = args.stoppable
         self.refreshable = args.refreshable

     async def handle_request(
         self, request: aiohttp.web.BaseRequest
     ) -> typing.Union[aiohttp.web.Response, aiohttp.web.StreamResponse]:
         """Generic handler for all incoming HTTP requests"""

         # Define response headers first...
         headers = {
             "Server": "Apache Pony Mail (Foal/%s ~%s)" % (PONYMAIL_FOAL_VERSION, PONYMAIL_SERVER_VERSION),
         }

         if self.api_logger:
             self.api_logger.info(request.raw_path)

         # Figure out who is going to handle this request, if any
         # We are backwards compatible with the old Lua interface URLs
         body_type = "form"
         # Support URLs of form /api/handler/extra?query
         parts = request.path.split("/")
         if len(parts) < 3:
             return aiohttp.web.Response(
                 headers=headers, status=404, text="API Endpoint not found!"
             )
         handler = parts[2]
         # handle test requests
         if self.stoppable and handler == 'stop':
             self.background_event.set()
             return aiohttp.web.Response(headers=headers, status=200, text='Stop requested\n')
         if self.refreshable and handler == 'refresh':
             await plugins.background.get_data(self)
             return aiohttp.web.Response(headers=headers, status=200, text='Refresh performed\n')

         if handler.endswith(".lua"):
             body_type = "form"
             handler = handler[:-4]
         if handler.endswith(".json"):
             body_type = "json"
             handler = handler[:-5]

         # Parse form data if any
         try:
             indata = await plugins.formdata.parse_formdata(body_type, request)
             if self.api_logger:
                 self.api_logger.info(indata)
         except ValueError as e:
             return aiohttp.web.Response(headers=headers, status=400, text=str(e))

         # Find a handler, or 404
         if handler in self.handlers:
             session = await plugins.session.get_session(self, request)
             try:
                 # Wait for endpoint response. This is typically JSON in case of success,
                 # but could be an exception (that needs a traceback) OR
                 # it could be a custom response, which we just pass along to the client.
                 xhandler = self.handlers[handler]
                 if isinstance(xhandler, plugins.server.StreamingEndpoint):
                     output = await xhandler.exec(self, request, session, indata)
                 elif isinstance(xhandler, plugins.server.Endpoint):
                     output = await xhandler.exec(self, session, indata)
                 if session.database:
                     self.dbpool.put_nowait(session.database)
                     self.dbpool.task_done()
                     session.database = None
                 if isinstance(output, aiohttp.web.Response) or isinstance(output, aiohttp.web.StreamResponse):
                     return output
                 if output:
                     jsout = await self.runners.run(json.dumps, output, indent=2)
                     headers["content-type"] = "application/json"
                     headers["Content-Length"] = str(len(jsout))
                     return aiohttp.web.Response(headers=headers, status=200, text=jsout)
                 return aiohttp.web.Response(
                     headers=headers, status=404, text="Content not found"
                 )
             # If a handler hit an exception, we need to print that exception somewhere,
             # either to the web client or stderr:
             except:
                 if session.database:
                     self.dbpool.put_nowait(session.database)
                     self.dbpool.task_done()
                     session.database = None
                 exc_type, exc_value, exc_traceback = sys.exc_info()
                 err = "\n".join(
                     traceback.format_exception(exc_type, exc_value, exc_traceback)
                 )
                 # By default, we print the traceback to the user, for easy debugging.
                 if self.config.ui.traceback:
                     return aiohttp.web.Response(
                         headers=headers, status=500, text="API error occurred: \n" + err
                     )
                 # If client traceback is disabled, we print it to stderr instead, but leave an
                 # error ID for the client to report back to the admin. Every line of the traceback
                 # will have this error ID at the beginning of the line, for easy grepping.
                 # We only need a short ID here, let's pick 18 chars.
                 eid = str(uuid.uuid4())[:18]
                 sys.stderr.write("API Endpoint %s got into trouble (%s): \n" % (request.path, eid))
                 for line in err.split("\n"):
                     sys.stderr.write("%s: %s\n" % (eid, line))
                 return aiohttp.web.Response(
                     headers=headers, status=500, text="API error occurred. The application journal will have "
                                                       "information. Error ID: %s" % eid
                 )
         else:
             return aiohttp.web.Response(
                 headers=headers, status=404, text="API Endpoint not found!"
             )

     async def server_loop(self):
         self.server = aiohttp.web.Server(self.handle_request)
         runner = aiohttp.web.ServerRunner(self.server)
         await runner.setup()
         site = aiohttp.web.TCPSite(
             runner, self.config.server.ip, self.config.server.port
         )
         await site.start()
         print(
             "==== Serving up Pony goodness at %s:%s ===="
             % (self.config.server.ip, self.config.server.port)
         )
         await plugins.background.run_tasks(self)
         await self.cleanup()
         await site.stop() # try to clean up

     async def cleanup(self):
         while not self.dbpool.empty():
             await self.dbpool.get_nowait().client.close()

     def run(self):
         # get_event_loop is deprecated in 3.10, but the replacment new_event_loop
         # does not seem to work properly in earlier versions
         if sys.version_info.minor < 10:
             loop = asyncio.get_event_loop()
         else:
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
         try:
             loop.run_until_complete(self.server_loop())
         except KeyboardInterrupt:
             self.background_event.set()
             loop.run_until_complete(self.cleanup())
         loop.close()


 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--config",
         help="Configuration file to load (default: ponymail.yaml)",
         default="ponymail.yaml",
     )
     parser.add_argument(
         "--logger",
         help="elasticsearch level (e.g. INFO or DEBUG)",
     )
     parser.add_argument(
         "--trace",
         help="elasticsearch.trace level (e.g. INFO or DEBUG)",
     )
     parser.add_argument(
         "--apilog",
         help="api log level (e.g. INFO or DEBUG)",
     )
     parser.add_argument(
         "--stoppable",
         action='store_true',
         help="Allow remote stop for testing",
     )
     parser.add_argument(
         "--refreshable",
         action='store_true',
         help="Allow remote refresh for testing",
     )
     parser.add_argument(
         "--testendpoints",
         action='store_true',
         help="Enable test endpoints",
     )
     cliargs = parser.parse_args()
     Server(cliargs).run()
	#!/usr/bin/env python3
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you 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.

	"""Apache Pony Mail, Codename Foal - A Python variant of Pony Mail"""
	import argparse
	import asyncio
	import importlib
	import json
	import os
	import sys
	from time import sleep
	import traceback
	import typing

	import aiohttp.web
	import yaml
	import uuid

	import plugins.background
	import plugins.configuration
	import plugins.database
	import plugins.formdata
	import plugins.offloader
	import plugins.server
	import plugins.session

	PONYMAIL_FOAL_VERSION = "0.1.0"
	from server_version import PONYMAIL_SERVER_VERSION


	# Certain environments such as MinGW-w64 will not register as a TTY and uses buffered output.
	# In such cases, we need to force a flush of each print, or nothing will show.
	if not sys.stdout.buffer.isatty():
	import functools
	print = functools.partial(print, flush=True)


	class Server(plugins.server.BaseServer):
	"""Main server class, responsible for handling requests and scheduling offloader threads """

	def _load_endpoint(self, subdir):
	for endpoint_file in sorted(os.listdir(subdir)):
	if endpoint_file.endswith(".py"):
	endpoint = endpoint_file[:-3]
	m = importlib.import_module(f"{subdir}.{endpoint}")
	if hasattr(m, "register"):
	self.handlers[endpoint] = m.__getattribute__("register")(self)
	print(f"Registered endpoint /api/{endpoint}")
	else:
	print(
	f"Could not find entry point 'register()' in {endpoint_file}, skipping!"
	)

	def __init__(self, args: argparse.Namespace):
	print(
	"==== Apache Pony Mail (Foal v/%s ~%s) starting... ====" % (PONYMAIL_FOAL_VERSION, PONYMAIL_SERVER_VERSION)
	)
	# Load configuration
	yml = yaml.safe_load(open(args.config))
	self.config = plugins.configuration.Configuration(yml)
	self.data = plugins.configuration.InterData()
	self.handlers = dict()
	self.dbpool = asyncio.Queue()
	self.runners = plugins.offloader.ExecutorPool()
	self.server = None
	self.streamlock = asyncio.Lock()
	self.api_logger = None
	self.foal_version = PONYMAIL_FOAL_VERSION
	self.server_version = PONYMAIL_SERVER_VERSION
	self.stoppable = False # allow remote stop for tests
	self.background_event = asyncio.Event() # for background task to wait on

	# Make a pool of database connections for async queries
	pool_size = self.config.database.pool_size
	if pool_size < 1:
	raise ValueError(f"pool_size {pool_size} must be > 0")
	for _ in range(0, pool_size): # stop value is exclusive
	self.dbpool.put_nowait(plugins.database.Database(self.config.database))

	# Load each URL endpoint
	if args.testendpoints:
	print(" Loading additional testing endpoints ")
	self._load_endpoint("testendpoints")
	print()
	self._load_endpoint("endpoints")

	if args.logger:
	import logging
	es_logger = logging.getLogger('elasticsearch')
	es_logger.setLevel(args.logger)
	es_logger.addHandler(logging.StreamHandler())
	if args.trace:
	import logging
	es_trace_logger = logging.getLogger('elasticsearch.trace')
	es_trace_logger.setLevel(args.trace)
	es_trace_logger.addHandler(logging.StreamHandler())
	if args.apilog:
	import logging
	self.api_logger = logging.getLogger('ponymail.apilog')
	self.api_logger.setLevel(args.apilog)
	self.api_logger.addHandler(logging.StreamHandler())
	self.stoppable = args.stoppable
	self.refreshable = args.refreshable

	async def handle_request(
	self, request: aiohttp.web.BaseRequest
	) -> typing.Union[aiohttp.web.Response, aiohttp.web.StreamResponse]:
	"""Generic handler for all incoming HTTP requests"""

	# Define response headers first...
	headers = {
	"Server": "Apache Pony Mail (Foal/%s ~%s)" % (PONYMAIL_FOAL_VERSION, PONYMAIL_SERVER_VERSION),
	}

	if self.api_logger:
	self.api_logger.info(request.raw_path)

	# Figure out who is going to handle this request, if any
	# We are backwards compatible with the old Lua interface URLs
	body_type = "form"
	# Support URLs of form /api/handler/extra?query
	parts = request.path.split("/")
	if len(parts) < 3:
	return aiohttp.web.Response(
	headers=headers, status=404, text="API Endpoint not found!"
	)
	handler = parts[2]
	# handle test requests
	if self.stoppable and handler == 'stop':
	self.background_event.set()
	return aiohttp.web.Response(headers=headers, status=200, text='Stop requested\n')
	if self.refreshable and handler == 'refresh':
	await plugins.background.get_data(self)
	return aiohttp.web.Response(headers=headers, status=200, text='Refresh performed\n')

	if handler.endswith(".lua"):
	body_type = "form"
	handler = handler[:-4]
	if handler.endswith(".json"):
	body_type = "json"
	handler = handler[:-5]

	# Parse form data if any
	try:
	indata = await plugins.formdata.parse_formdata(body_type, request)
	if self.api_logger:
	self.api_logger.info(indata)
	except ValueError as e:
	return aiohttp.web.Response(headers=headers, status=400, text=str(e))

	# Find a handler, or 404
	if handler in self.handlers:
	session = await plugins.session.get_session(self, request)
	try:
	# Wait for endpoint response. This is typically JSON in case of success,
	# but could be an exception (that needs a traceback) OR
	# it could be a custom response, which we just pass along to the client.
	xhandler = self.handlers[handler]
	if isinstance(xhandler, plugins.server.StreamingEndpoint):
	output = await xhandler.exec(self, request, session, indata)
	elif isinstance(xhandler, plugins.server.Endpoint):
	output = await xhandler.exec(self, session, indata)
	if session.database:
	self.dbpool.put_nowait(session.database)
	self.dbpool.task_done()
	session.database = None
	if isinstance(output, aiohttp.web.Response) or isinstance(output, aiohttp.web.StreamResponse):
	return output
	if output:
	jsout = await self.runners.run(json.dumps, output, indent=2)
	headers["content-type"] = "application/json"
	headers["Content-Length"] = str(len(jsout))
	return aiohttp.web.Response(headers=headers, status=200, text=jsout)
	return aiohttp.web.Response(
	headers=headers, status=404, text="Content not found"
	)
	# If a handler hit an exception, we need to print that exception somewhere,
	# either to the web client or stderr:
	except:
	if session.database:
	self.dbpool.put_nowait(session.database)
	self.dbpool.task_done()
	session.database = None
	exc_type, exc_value, exc_traceback = sys.exc_info()
	err = "\n".join(
	traceback.format_exception(exc_type, exc_value, exc_traceback)
	)
	# By default, we print the traceback to the user, for easy debugging.
	if self.config.ui.traceback:
	return aiohttp.web.Response(
	headers=headers, status=500, text="API error occurred: \n" + err
	)
	# If client traceback is disabled, we print it to stderr instead, but leave an
	# error ID for the client to report back to the admin. Every line of the traceback
	# will have this error ID at the beginning of the line, for easy grepping.
	# We only need a short ID here, let's pick 18 chars.
	eid = str(uuid.uuid4())[:18]
	sys.stderr.write("API Endpoint %s got into trouble (%s): \n" % (request.path, eid))
	for line in err.split("\n"):
	sys.stderr.write("%s: %s\n" % (eid, line))
	return aiohttp.web.Response(
	headers=headers, status=500, text="API error occurred. The application journal will have "
	"information. Error ID: %s" % eid
	)
	else:
	return aiohttp.web.Response(
	headers=headers, status=404, text="API Endpoint not found!"
	)

	async def server_loop(self):
	self.server = aiohttp.web.Server(self.handle_request)
	runner = aiohttp.web.ServerRunner(self.server)
	await runner.setup()
	site = aiohttp.web.TCPSite(
	runner, self.config.server.ip, self.config.server.port
	)
	await site.start()
	print(
	"==== Serving up Pony goodness at %s:%s ===="
	% (self.config.server.ip, self.config.server.port)
	)
	await plugins.background.run_tasks(self)
	await self.cleanup()
	await site.stop() # try to clean up

	async def cleanup(self):
	while not self.dbpool.empty():
	await self.dbpool.get_nowait().client.close()

	def run(self):
	# get_event_loop is deprecated in 3.10, but the replacment new_event_loop
	# does not seem to work properly in earlier versions
	if sys.version_info.minor < 10:
	loop = asyncio.get_event_loop()
	else:
	loop = asyncio.new_event_loop()
	asyncio.set_event_loop(loop)
	try:
	loop.run_until_complete(self.server_loop())
	except KeyboardInterrupt:
	self.background_event.set()
	loop.run_until_complete(self.cleanup())
	loop.close()


	if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument(
	"--config",
	help="Configuration file to load (default: ponymail.yaml)",
	default="ponymail.yaml",
	)
	parser.add_argument(
	"--logger",
	help="elasticsearch level (e.g. INFO or DEBUG)",
	)
	parser.add_argument(
	"--trace",
	help="elasticsearch.trace level (e.g. INFO or DEBUG)",
	)
	parser.add_argument(
	"--apilog",
	help="api log level (e.g. INFO or DEBUG)",
	)
	parser.add_argument(
	"--stoppable",
	action='store_true',
	help="Allow remote stop for testing",
	)
	parser.add_argument(
	"--refreshable",
	action='store_true',
	help="Allow remote refresh for testing",
	)
	parser.add_argument(
	"--testendpoints",
	action='store_true',
	help="Enable test endpoints",
	)
	cliargs = parser.parse_args()
	Server(cliargs).run()