<commit_before># -*- coding: utf-8 -*-
"""
Created on Mon Oct 24 14:07:10 2016

@author: jonas.guetter
"""

import networkx as nx
from sqlalchemy.sql import text
from oemof import db
from shapely.wkt import loads
import math
import time

t0 = time.clock()

# connect to database
engine = db.engine(section='oedb')
conn = engine.connect()

# Truncate result table
#s = text("Truncate TABLE model_draft.ego_grid_lv_onts")
#result = conn.execute(s)
#result.close()

s = text("SELECT COUNT(id) AS numberofloadareas FROM calc_ego_loads.ego_deu_load_area")
result = conn.execute(s)
for row in result:
    numberofloadareas = row['numberofloadareas']
result.close()

# Iterate over load areas
for w in range(200000,numberofloadareas):

    # Initialize variables
    onts = []
    distmax = 250
    xcoord = []
    ycoord = []
    dominated = []
    ont = []
    street_geoms = []
    street_gids = []

    print ('new iteration step begins  ', (time.clock() - t0, "seconds process time"))

    # Perform query
    s = text(
            "SELECT ST_AsText(ST_UNION(ST_MULTI(geom))) AS geom, ST_GEOMETRYTYPE(ST_UNION(ST_MULTI(geom))) AS type, \
            AVG(ontnumber) AS ontnumber FROM model_draft.ego_grid_lv_streets WHERE load_area_id = :x")
    result = conn.execute(s, x = w)
    print ("load_area_id:",w)
#    print("type:"+str(type(result)))
    # Get number of ont for this load area
    for row in result:
        #print("row:"+str(row))
        numberofonts = row['ontnumber']
        result2 = row['geom']
        resulttype = row['type']

    # Do nothing if there is no or only one street in the load area:
    if (str(result2) == "None"):
        print ("passed")
        pass
    elif (resulttype == "ST_LineString"):
        print('only one street in the load area')
    else:

        #print ("result2:" + str(result2))
        wkt = loads(result2)
        print ('Create graph  ', (time.clock() - t0, "seconds process time"))
        # Create graph from lines
        graph = nx.Graph()
        for line in wkt:
            for seg_start, seg_end in zip(list(line.coords),list(line.coords)[1:]):
                graph.add_edge(seg_start, seg_end,)

        result.close()

        print ('change to undirected graph  ', (time.clock() - t0, "seconds process time"))
        # change to undirected graph
        graph = graph.to_undirected()

        print ('number of nodes:')
        print (graph.number_of_nodes())

        # Leave out very big load_areas due to performance issues
        if(graph.number_of_nodes() > 9000):
            print ("Street network is too large")
        else:
            # Perform positioning algorithm on graph

            # Step 1: Mark every node as undominated
            print ('Step 1  ', (time.clock() - t0, "seconds process time"))
            for x in range(0, graph.number_of_nodes()):
                graph.node[graph.nodes()[x]]['dominated'] = None

            # Step 2
            print ('Step 2.1  ', (time.clock() - t0, "seconds process time"))
            # Step 2.1: Count the number of adjacent nodes for each node:
            for x in range(0, graph.number_of_nodes()):
                adjacent_nodes = len(graph.neighbors(graph.nodes()[x]))
                graph.node[graph.nodes()[x]]['adjacent_undominated'] = adjacent_nodes
                graph.node[graph.nodes()[x]]['ONT'] = False

            # Step 2.2: Calculate the length of each edge
            print ('Step 2.2  ', (time.clock() - t0, "seconds process time"))
            for x in range(0, graph.number_of_edges()):
                first = graph.edges(data=True)[x]
                distance = math.sqrt(sum([(a - b) ** 2 for a, b in zip(first[0],first[1])]))
                graph.edge[first[0]][first[1]]['weight'] = distance
            #print ('nodes:')
            #print (graph.nodes(data=True))

            print ('Step 2.3  ', (time.clock() - t0, "seconds process time"))
            for x in range(0, int(numberofonts)):
                # Step 2.3: Add the node with the most adjacent undominated nodes to onts
                max_neighbors = max(nx.get_node_attributes(graph, 'adjacent_undominated').values())
                #print('max neighbors:' + str(max_neighbors))
                for y in range (0, graph.number_of_nodes()):
                    if (graph.node[graph.nodes()[y]]['adjacent_undominated'] == max_neighbors):
                        nextont = graph.nodes()[y]
                        graph.node[graph.nodes()[y]]['ONT'] = True
                        onts.append (nextont)
                        break

                # Step 3: Mark every node within the range of the previously added ont as dominated

                # Step 3.1: Calculate the distance for each node and the previously added ont
                distances = nx.single_source_dijkstra_path_length (graph, source =
                   nextont, cutoff=distmax, weight= 'weight')

                # Step 3.2: Mark every node within the range of the previously added ont as dominated
                #print ('distances:')
                #print (distances)
                for y in distances.keys():
                    if (graph.node[y]['dominated'] == None):
                        graph.node[y]['dominated'] = x

                        # Step 3.3: For every neighbor of those nodes: Decrease number of adjacent undominated nodes by 1
                        for z in graph.neighbors(y):
                            graph.node[z]['adjacent_undominated'] -= 1
                    else:

                        # if a node is already dominated, change the dominating ont if the distance to the new ont is shorter
                        ont_new = nextont
                        ont_old = onts[graph.node[y]['dominated']]
                        #print ("ont_old:")
                        #print (ont_old)
                        point = graph.nodes()[graph.nodes().index(y)]
                        pathlength_old = nx.shortest_path_length(graph,source = point,target = ont_old, weight = 'weight')
                        pathlength_new = nx.shortest_path_length(graph,source = point,target = ont_new, weight = 'weight')
                        if ( pathlength_new <= pathlength_old ):
                            graph.node[y]['dominated'] = x

        # assign every node which is not already dominated to a ont
#        print ('Assign rest nodes  ', (time.clock() - t0, "seconds process time"))
#        for x in range (0,graph.number_of_nodes()):
#            if (graph.node[graph.nodes()[x]]['dominated'] == None):
#                # Calculate the distance to each ont
#                pathlengths = []
#                for y in range(0, int(numberofonts)):
#                    # Check if path between node and ont exists
#                    if (nx.has_path(graph,source = graph.nodes()[x], target = onts[y])):
#                        pathlengths.append(nx.shortest_path_length(graph,source = graph.nodes()[x], target = onts[y], weight = 'weight'))
#                    else:
#                        pathlengths.append(999999999999999)
#                # Assign the ont with the minimal pathlength to the node
#                if (min(pathlengths) < 999999999999999):
#                    graph.node[graph.nodes()[x]]['dominated'] = pathlengths.index(min(pathlengths))

        # Assign onts to edges
#        print ('Assign onts to edges  ', (time.clock() - t0, "seconds process time"))
#        for x in range(0, graph.number_of_edges()):
#            ed = graph.edges(data=True)[x]
#            graph.edge[ed[0]][ed[1]]['dominated'] = graph.node[ed[1]]['dominated']

        # create output graph
        out_graph = graph.copy()
        out_graph.remove_nodes_from(graph.nodes())
        out_graph.add_nodes_from(onts)

        # Insert values into table
        s = text(
                "Insert into model_draft.ego_grid_lv_onts (geom) SELECT ST_SetSRID(ST_MAKEPOINT(:x,:y),3035)")

        for z in range(0, out_graph.number_of_nodes()):
            #print (out_graph.nodes()[z][0])
            result = conn.execute(s, x = out_graph.nodes()[z][0],y = out_graph.nodes()[z][1])

        result.close()

# Add MV substation information to onts
s = text(
    "ALTER TABLE model_draft.ego_grid_mvlv_onts\
    ADD COLUMN subst_id integer;\
    \
    UPDATE model_draft.ego_grid_mvlv_onts AS t1\
    SET subst_id = t2.subst_id FROM\
        (SELECT onts.id AS id, MAX(district.subst_id) AS subst_id\
        FROM model_draft.ego_grid_mvlv_onts AS onts, calc_ego_grid_district.grid_district AS district\
        WHERE ST_INTERSECTS (onts.geom,district.geom)\
        GROUP BY onts.id) AS t2\
    WHERE t1.id = t2.id"
        )

result = conn.execute(s)
result.close()

# Add load area information to onts
s = text(
    "ALTER TABLE model_draft.ego_grid_mvlv_onts\
    ADD COLUMN la_id integer;\
    \
    UPDATE model_draft.ego_grid_mvlv_onts AS t1\
    SET la_id = t2.la_id FROM\
        (SELECT onts.id AS id, MAX(larea.id) AS la_id\
        FROM model_draft.ego_grid_mvlv_onts AS onts, calc_ego_loads.ego_deu_load_area AS larea\
        WHERE ST_INTERSECTS (onts.geom,larea.geom)\
        GROUP BY onts.id) AS t2\
    WHERE t1.id = t2.id"
        )

result = conn.execute(s)
result.close()

# Close db connection
conn.close()

print ((time.clock() - t0)/3600, "hours process time")

## Store results ins lists
#for x in range(0, graph.number_of_nodes()):
#    ont.append(graph.node[graph.nodes()[x]]['ONT'])
#    dominated.append(graph.node[graph.nodes()[x]]['dominated'])
#    xc,yc = graph.nodes()[x]
#    xcoord.append(xc)
#    ycoord.append(yc)
#
#
## write results into csv file
#csv_name = 'process_eGo_lv_grid.csv'
#with open(csv_name,'w',newline = '') as foo:
#    writer = csv.writer(foo,delimiter = ',')
#    for x in range (0, graph.number_of_nodes()):
#        writer.writerow([xcoord[x],ycoord[x],dominated[x], ont[x]])
#
#
#
#

<commit_msg>add license header to process_eGo_lv_grid_multiple.py<commit_after>"""
Create LV grid districts
"""

__copyright__ 	= "Reiner Lemoine Institut gGmbH"
__license__ 	= "GNU Affero General Public License Version 3 (AGPL-3.0)"
__url__         = "https://github.com/openego/data_processing/blob/master/LICENSE"
__author__ = "jong42"

import networkx as nx
from sqlalchemy.sql import text
from oemof import db
from shapely.wkt import loads
import math
import time

t0 = time.clock()

# connect to database
engine = db.engine(section='oedb')
conn = engine.connect()

# Truncate result table
#s = text("Truncate TABLE model_draft.ego_grid_lv_onts")
#result = conn.execute(s)
#result.close()

s = text("SELECT COUNT(id) AS numberofloadareas FROM calc_ego_loads.ego_deu_load_area")
result = conn.execute(s)
for row in result:
    numberofloadareas = row['numberofloadareas']
result.close()

# Iterate over load areas
for w in range(200000,numberofloadareas):

    # Initialize variables
    onts = []
    distmax = 250
    xcoord = []
    ycoord = []
    dominated = []
    ont = []
    street_geoms = []
    street_gids = []

    print ('new iteration step begins  ', (time.clock() - t0, "seconds process time"))

    # Perform query
    s = text(
            "SELECT ST_AsText(ST_UNION(ST_MULTI(geom))) AS geom, ST_GEOMETRYTYPE(ST_UNION(ST_MULTI(geom))) AS type, \
            AVG(ontnumber) AS ontnumber FROM model_draft.ego_grid_lv_streets WHERE load_area_id = :x")
    result = conn.execute(s, x = w)
    print ("load_area_id:",w)
#    print("type:"+str(type(result)))
    # Get number of ont for this load area
    for row in result:
        #print("row:"+str(row))
        numberofonts = row['ontnumber']
        result2 = row['geom']
        resulttype = row['type']

    # Do nothing if there is no or only one street in the load area:
    if (str(result2) == "None"):
        print ("passed")
        pass
    elif (resulttype == "ST_LineString"):
        print('only one street in the load area')
    else:

        #print ("result2:" + str(result2))
        wkt = loads(result2)
        print ('Create graph  ', (time.clock() - t0, "seconds process time"))
        # Create graph from lines
        graph = nx.Graph()
        for line in wkt:
            for seg_start, seg_end in zip(list(line.coords),list(line.coords)[1:]):
                graph.add_edge(seg_start, seg_end,)

        result.close()

        print ('change to undirected graph  ', (time.clock() - t0, "seconds process time"))
        # change to undirected graph
        graph = graph.to_undirected()

        print ('number of nodes:')
        print (graph.number_of_nodes())

        # Leave out very big load_areas due to performance issues
        if(graph.number_of_nodes() > 9000):
            print ("Street network is too large")
        else:
            # Perform positioning algorithm on graph

            # Step 1: Mark every node as undominated
            print ('Step 1  ', (time.clock() - t0, "seconds process time"))
            for x in range(0, graph.number_of_nodes()):
                graph.node[graph.nodes()[x]]['dominated'] = None

            # Step 2
            print ('Step 2.1  ', (time.clock() - t0, "seconds process time"))
            # Step 2.1: Count the number of adjacent nodes for each node:
            for x in range(0, graph.number_of_nodes()):
                adjacent_nodes = len(graph.neighbors(graph.nodes()[x]))
                graph.node[graph.nodes()[x]]['adjacent_undominated'] = adjacent_nodes
                graph.node[graph.nodes()[x]]['ONT'] = False

            # Step 2.2: Calculate the length of each edge
            print ('Step 2.2  ', (time.clock() - t0, "seconds process time"))
            for x in range(0, graph.number_of_edges()):
                first = graph.edges(data=True)[x]
                distance = math.sqrt(sum([(a - b) ** 2 for a, b in zip(first[0],first[1])]))
                graph.edge[first[0]][first[1]]['weight'] = distance
            #print ('nodes:')
            #print (graph.nodes(data=True))

            print ('Step 2.3  ', (time.clock() - t0, "seconds process time"))
            for x in range(0, int(numberofonts)):
                # Step 2.3: Add the node with the most adjacent undominated nodes to onts
                max_neighbors = max(nx.get_node_attributes(graph, 'adjacent_undominated').values())
                #print('max neighbors:' + str(max_neighbors))
                for y in range (0, graph.number_of_nodes()):
                    if (graph.node[graph.nodes()[y]]['adjacent_undominated'] == max_neighbors):
                        nextont = graph.nodes()[y]
                        graph.node[graph.nodes()[y]]['ONT'] = True
                        onts.append (nextont)
                        break

                # Step 3: Mark every node within the range of the previously added ont as dominated

                # Step 3.1: Calculate the distance for each node and the previously added ont
                distances = nx.single_source_dijkstra_path_length (graph, source =
                   nextont, cutoff=distmax, weight= 'weight')

                # Step 3.2: Mark every node within the range of the previously added ont as dominated
                #print ('distances:')
                #print (distances)
                for y in distances.keys():
                    if (graph.node[y]['dominated'] == None):
                        graph.node[y]['dominated'] = x

                        # Step 3.3: For every neighbor of those nodes: Decrease number of adjacent undominated nodes by 1
                        for z in graph.neighbors(y):
                            graph.node[z]['adjacent_undominated'] -= 1
                    else:

                        # if a node is already dominated, change the dominating ont if the distance to the new ont is shorter
                        ont_new = nextont
                        ont_old = onts[graph.node[y]['dominated']]
                        #print ("ont_old:")
                        #print (ont_old)
                        point = graph.nodes()[graph.nodes().index(y)]
                        pathlength_old = nx.shortest_path_length(graph,source = point,target = ont_old, weight = 'weight')
                        pathlength_new = nx.shortest_path_length(graph,source = point,target = ont_new, weight = 'weight')
                        if ( pathlength_new <= pathlength_old ):
                            graph.node[y]['dominated'] = x

        # assign every node which is not already dominated to a ont
#        print ('Assign rest nodes  ', (time.clock() - t0, "seconds process time"))
#        for x in range (0,graph.number_of_nodes()):
#            if (graph.node[graph.nodes()[x]]['dominated'] == None):
#                # Calculate the distance to each ont
#                pathlengths = []
#                for y in range(0, int(numberofonts)):
#                    # Check if path between node and ont exists
#                    if (nx.has_path(graph,source = graph.nodes()[x], target = onts[y])):
#                        pathlengths.append(nx.shortest_path_length(graph,source = graph.nodes()[x], target = onts[y], weight = 'weight'))
#                    else:
#                        pathlengths.append(999999999999999)
#                # Assign the ont with the minimal pathlength to the node
#                if (min(pathlengths) < 999999999999999):
#                    graph.node[graph.nodes()[x]]['dominated'] = pathlengths.index(min(pathlengths))

        # Assign onts to edges
#        print ('Assign onts to edges  ', (time.clock() - t0, "seconds process time"))
#        for x in range(0, graph.number_of_edges()):
#            ed = graph.edges(data=True)[x]
#            graph.edge[ed[0]][ed[1]]['dominated'] = graph.node[ed[1]]['dominated']

        # create output graph
        out_graph = graph.copy()
        out_graph.remove_nodes_from(graph.nodes())
        out_graph.add_nodes_from(onts)

        # Insert values into table
        s = text(
                "Insert into model_draft.ego_grid_lv_onts (geom) SELECT ST_SetSRID(ST_MAKEPOINT(:x,:y),3035)")

        for z in range(0, out_graph.number_of_nodes()):
            #print (out_graph.nodes()[z][0])
            result = conn.execute(s, x = out_graph.nodes()[z][0],y = out_graph.nodes()[z][1])

        result.close()

# Add MV substation information to onts
s = text(
    "ALTER TABLE model_draft.ego_grid_mvlv_onts\
    ADD COLUMN subst_id integer;\
    \
    UPDATE model_draft.ego_grid_mvlv_onts AS t1\
    SET subst_id = t2.subst_id FROM\
        (SELECT onts.id AS id, MAX(district.subst_id) AS subst_id\
        FROM model_draft.ego_grid_mvlv_onts AS onts, calc_ego_grid_district.grid_district AS district\
        WHERE ST_INTERSECTS (onts.geom,district.geom)\
        GROUP BY onts.id) AS t2\
    WHERE t1.id = t2.id"
        )

result = conn.execute(s)
result.close()

# Add load area information to onts
s = text(
    "ALTER TABLE model_draft.ego_grid_mvlv_onts\
    ADD COLUMN la_id integer;\
    \
    UPDATE model_draft.ego_grid_mvlv_onts AS t1\
    SET la_id = t2.la_id FROM\
        (SELECT onts.id AS id, MAX(larea.id) AS la_id\
        FROM model_draft.ego_grid_mvlv_onts AS onts, calc_ego_loads.ego_deu_load_area AS larea\
        WHERE ST_INTERSECTS (onts.geom,larea.geom)\
        GROUP BY onts.id) AS t2\
    WHERE t1.id = t2.id"
        )

result = conn.execute(s)
result.close()

# Close db connection
conn.close()

print ((time.clock() - t0)/3600, "hours process time")

## Store results ins lists
#for x in range(0, graph.number_of_nodes()):
#    ont.append(graph.node[graph.nodes()[x]]['ONT'])
#    dominated.append(graph.node[graph.nodes()[x]]['dominated'])
#    xc,yc = graph.nodes()[x]
#    xcoord.append(xc)
#    ycoord.append(yc)
#
#
## write results into csv file
#csv_name = 'process_eGo_lv_grid.csv'
#with open(csv_name,'w',newline = '') as foo:
#    writer = csv.writer(foo,delimiter = ',')
#    for x in range (0, graph.number_of_nodes()):
#        writer.writerow([xcoord[x],ycoord[x],dominated[x], ont[x]])
#
#
#
#

<|endoftext|>


<commit_before>import aiofiles
import json
import os
import subprocess
import tempfile

import virtool.errors
import virtool.github
import virtool.updates
import virtool.utils

PROJECTION = [
    "_id",
    "cluster",
    "names",
    "count",
    "families"
]

LATEST_RELEASE_URL = "https://api.github.com/repos/virtool/virtool-hmm/releases/latest"

async def hmmstat(loop, path):
    """

    :param loop: the application loop
    :param executor: the application executor
    :param path: the path to the profiles.hmm file

    :return: a list of profiles in the file
    :rtype: Coroutine[list]

    """
    if not os.path.isfile(path):
        raise FileNotFoundError("HMM file does not exist")

    command = ["hmmstat", path]

    output = await loop.run_in_executor(None, subprocess.check_output, command)

    result = [line.split() for line in output.decode("utf-8").split("\n") if line and line[0] != "#"]

    return [{
        "cluster": int(line[1].replace("vFam_", "")),
        "count": int(line[3]),
        "length": int(line[5])
    } for line in result]

async def get_referenced_hmm_ids(db):
    """
    Returns a list of HMM document ids that are referenced in analyses documents.

    :param db: the application database client
    :type db: :class:`~motor.motor_asyncio.AsyncIOMotorClient`

    :return: HMM document ids
    :rtype: Coroutine[list]

    """
    agg = await db.analyses.aggregate([
        {"$match": {
            "algorithm": "nuvs"
        }},
        {"$project": {
            "results.orfs.hits.hit": True
        }},
        {"$unwind": "$results"},
        {"$unwind": "$results.orfs"},
        {"$unwind": "$results.orfs.hits"},
        {"$group": {
            "_id": "$results.orfs.hits.hit"
        }}
    ]).to_list(None)

    return list(set(a["_id"] for a in agg))

async def update_process(db, dispatch, progress, step=None, error=None):
    return await virtool.utils.update_status_process(db, dispatch, "hmm_install", progress, step)

async def get_asset(server_version, username, token):
    """
    Get the asset information associated with the latest HMM release.

    :param server_version: the current server version
    :type server_version: str

    :param username: the GitHub username to use for auth
    :type username: str

    :param token: the GitHub token to use for auth
    :type  token: str

    :return: the assets, a dict containing the download url and size in a tuple
    :rtype: Coroutine[dict]

    """
    release = await virtool.github.get(LATEST_RELEASE_URL, server_version, username, token)

    assets = list()

    for asset in release["assets"]:
        assets.append((asset["browser_download_url"], asset["size"]))

    return assets

async def install_official(loop, db, settings, dispatch, server_version, username=None, token=None):
    """
    Runs a background Task that:

        - downloads the official profiles.hmm.gz file
        - decompresses the profiles.hmm.gz file
        - moves the profiles.hmm file to the correct data path
        - downloads the official annotations.json.gz file
        - imports the annotations into the database

    Task reports the following stages to the hmm_install status document:

        1. check_github
        2. download
        3. decompress
        4. install_profiles
        5. import_annotations

    :param loop: the application event loop

    :param db: the application database client
    :type db: :class:`~motor.motor_asyncio.AsyncIOMotorClient`

    :param settings: the application Settings object
    :type settings: :class:`virtool.app_settings.Settings`

    :param dispatch: the application dispatch method
    :type dispatch: func

    :param server_version: the current server version
    :type server_version: str

    :param username: the GitHub username to use for auth
    :type username: str

    :param token: the GitHub token to use for auth
    :type  token: str

    """
    await update_process(db, dispatch, 0, step="check_github")

    assets = await get_asset(server_version, username, token)

    await db.status.update_one({"_id": "hmm_install"}, {
        "$set": {
            "download_size": int(assets[0][1])
        }
    })

    await update_process(db, dispatch, 0.5)

    if len(assets) == 0:
        # Stop the install process if one of the annotation or profile assets are not found.
        return await update_process(db, dispatch, 1.0, error="Missing HMM asset file")

    async def handler(progress):
        await update_process(db, dispatch, progress)

    with tempfile.TemporaryDirectory() as temp_path:
        target_path = os.path.join(temp_path, "vthmm.tar.gz")

        url, size = assets[0]

        await update_process(db, dispatch, 0, step="download")
        await virtool.github.download_asset(url, size, target_path, progress_handler=handler)

        await update_process(db, dispatch, 0, step="decompress")

        await loop.run_in_executor(None, virtool.github.decompress_asset_file, target_path, temp_path)

        await update_process(db, dispatch, 0, step="install_profiles")

        decompressed_path = os.path.join(temp_path, "hmm")
        install_path = os.path.join(settings.get("data_path"), "hmm", "profiles.hmm")
        await loop.run_in_executor(None, os.rename, os.path.join(decompressed_path, "profiles.hmm"), install_path)

        await update_process(db, dispatch, 0, step="import_annotations")

        async with aiofiles.open(os.path.join(decompressed_path, "annotations.json"), "r") as f:
            annotations = json.loads(await f.read())

        await insert_annotations(db, annotations, handler)

        await db.status.update_one({"_id": "hmm_install"}, {
            "$set": {
                "ready": True
            }
        })

        await update_process(db, dispatch, 1.0)

async def insert_annotations(db, annotations, progress_handler=None):
    existing_ids = set()

    count = 0
    total_count = len(annotations)

    for i in range(0, total_count, 30):
        chunk = annotations[i:i + 30]

        for annotation in chunk:
            _id = virtool.utils.random_alphanumeric(8, excluded=existing_ids)
            existing_ids.add(_id)
            annotation["_id"] = _id

        await db.hmm.insert_many(chunk)

        if progress_handler:
            count += len(chunk)
            await progress_handler(count / total_count)

<commit_msg>Delete unreferenced hmm annotations when installing from gh<commit_after>import aiofiles
import json
import os
import subprocess
import tempfile

import virtool.errors
import virtool.github
import virtool.updates
import virtool.utils

PROJECTION = [
    "_id",
    "cluster",
    "names",
    "count",
    "families"
]

LATEST_RELEASE_URL = "https://api.github.com/repos/virtool/virtool-hmm/releases/latest"

async def hmmstat(loop, path):
    """

    :param loop: the application loop
    :param executor: the application executor
    :param path: the path to the profiles.hmm file

    :return: a list of profiles in the file
    :rtype: Coroutine[list]

    """
    if not os.path.isfile(path):
        raise FileNotFoundError("HMM file does not exist")

    command = ["hmmstat", path]

    output = await loop.run_in_executor(None, subprocess.check_output, command)

    result = [line.split() for line in output.decode("utf-8").split("\n") if line and line[0] != "#"]

    return [{
        "cluster": int(line[1].replace("vFam_", "")),
        "count": int(line[3]),
        "length": int(line[5])
    } for line in result]

async def get_referenced_hmm_ids(db):
    """
    Returns a list of HMM document ids that are referenced in analyses documents.

    :param db: the application database client
    :type db: :class:`~motor.motor_asyncio.AsyncIOMotorClient`

    :return: HMM document ids
    :rtype: Coroutine[list]

    """
    agg = await db.analyses.aggregate([
        {"$match": {
            "algorithm": "nuvs"
        }},
        {"$project": {
            "results.orfs.hits.hit": True
        }},
        {"$unwind": "$results"},
        {"$unwind": "$results.orfs"},
        {"$unwind": "$results.orfs.hits"},
        {"$group": {
            "_id": "$results.orfs.hits.hit"
        }}
    ]).to_list(None)

    return list(set(a["_id"] for a in agg))

async def update_process(db, dispatch, progress, step=None, error=None):
    return await virtool.utils.update_status_process(db, dispatch, "hmm_install", progress, step)

async def get_asset(server_version, username, token):
    """
    Get the asset information associated with the latest HMM release.

    :param server_version: the current server version
    :type server_version: str

    :param username: the GitHub username to use for auth
    :type username: str

    :param token: the GitHub token to use for auth
    :type  token: str

    :return: the assets, a dict containing the download url and size in a tuple
    :rtype: Coroutine[dict]

    """
    release = await virtool.github.get(LATEST_RELEASE_URL, server_version, username, token)

    assets = list()

    for asset in release["assets"]:
        assets.append((asset["browser_download_url"], asset["size"]))

    return assets

async def install_official(loop, db, settings, dispatch, server_version, username=None, token=None):
    """
    Runs a background Task that:

        - downloads the official profiles.hmm.gz file
        - decompresses the profiles.hmm.gz file
        - moves the profiles.hmm file to the correct data path
        - downloads the official annotations.json.gz file
        - imports the annotations into the database

    Task reports the following stages to the hmm_install status document:

        1. check_github
        2. download
        3. decompress
        4. install_profiles
        5. import_annotations

    :param loop: the application event loop

    :param db: the application database client
    :type db: :class:`~motor.motor_asyncio.AsyncIOMotorClient`

    :param settings: the application Settings object
    :type settings: :class:`virtool.app_settings.Settings`

    :param dispatch: the application dispatch method
    :type dispatch: func

    :param server_version: the current server version
    :type server_version: str

    :param username: the GitHub username to use for auth
    :type username: str

    :param token: the GitHub token to use for auth
    :type  token: str

    """
    await update_process(db, dispatch, 0, step="check_github")

    assets = await get_asset(server_version, username, token)

    await db.status.update_one({"_id": "hmm_install"}, {
        "$set": {
            "download_size": int(assets[0][1])
        }
    })

    await update_process(db, dispatch, 0.5)

    if len(assets) == 0:
        # Stop the install process if one of the annotation or profile assets are not found.
        return await update_process(db, dispatch, 1.0, error="Missing HMM asset file")

    async def handler(progress):
        await update_process(db, dispatch, progress)

    with tempfile.TemporaryDirectory() as temp_path:
        target_path = os.path.join(temp_path, "vthmm.tar.gz")

        url, size = assets[0]

        await update_process(db, dispatch, 0, step="download")
        await virtool.github.download_asset(url, size, target_path, progress_handler=handler)

        await update_process(db, dispatch, 0, step="decompress")

        await loop.run_in_executor(None, virtool.github.decompress_asset_file, target_path, temp_path)

        await update_process(db, dispatch, 0, step="install_profiles")

        decompressed_path = os.path.join(temp_path, "hmm")
        install_path = os.path.join(settings.get("data_path"), "hmm", "profiles.hmm")
        await loop.run_in_executor(None, os.rename, os.path.join(decompressed_path, "profiles.hmm"), install_path)

        await update_process(db, dispatch, 0, step="import_annotations")

        async with aiofiles.open(os.path.join(decompressed_path, "annotations.json"), "r") as f:
            annotations = json.loads(await f.read())

        await insert_annotations(db, annotations, handler)

        await db.status.update_one({"_id": "hmm_install"}, {
            "$set": {
                "ready": True
            }
        })

        await update_process(db, dispatch, 1.0)

async def insert_annotations(db, annotations, progress_handler=None):
    referenced_ids = await get_referenced_hmm_ids(db)

    await db.hmm.delete_many({"_id": {"$nin": referenced_ids}})

    await db.hmm.update_many({}, {
        "$set": {
            "hidden": True
        }
    })

    existing_ids = set(await db.hmm.distinct("_id"))

    count = 0
    total_count = len(annotations)

    for i in range(0, total_count, 30):
        chunk = annotations[i:i + 30]

        for annotation in chunk:
            _id = virtool.utils.random_alphanumeric(8, excluded=existing_ids)
            existing_ids.add(_id)
            annotation.update({
                "_id": _id,
                "hidden": False
            })

        await db.hmm.insert_many(chunk)

        if progress_handler:
            count += len(chunk)
            await progress_handler(count / total_count)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""MgShip CLI classes."""
from __future__ import absolute_import

import click
import logging

from mgship import mgship
from mgship.destination import csv
from mgship.cliparam import DateTime, Loglevel
from mgship.util import utctimestamp, validate_past
from mgship.log import logger

@click.group()
@click.option('--log-level', default=None, type=Loglevel())
def main(log_level=None):
    """Console script for mgship"""
    if log_level is not None:
        logging.basicConfig(level=log_level)

def to_timestamp(ctx, param, value):
    if value is None:
        return None
    try:
        value = utctimestamp(value)
        if ctx.info_name == 'archive':
            validate_past(value)
        return value
    except (TypeError, ValueError) as e:
        logger.debug("bad timestamp parameter", exc_info=True)
        raise click.BadParameter(unicode(e))

@main.command()
@click.option('--begin', default=None, type=DateTime(),
              callback=to_timestamp,
              help='when to start archiving, as unix timestamp')
def archive(*args, **kwargs):
    mgship.Archive(csv.Destination(), *args, **kwargs).ship()

if __name__ == "__main__":
    main()
<commit_msg>enable both formats<commit_after># -*- coding: utf-8 -*-
"""MgShip CLI classes."""
from __future__ import absolute_import

import click
import logging

from mgship import mgship
from mgship.destination import csv, json
from mgship.cliparam import DateTime, Loglevel
from mgship.util import utctimestamp, validate_past
from mgship.log import logger

@click.group()
@click.option('--log-level', default=None, type=Loglevel())
def main(log_level=None):
    """Console script for mgship"""
    if log_level is not None:
        logging.basicConfig(level=log_level)

def to_timestamp(ctx, param, value):
    if value is None:
        return None
    try:
        value = utctimestamp(value)
        if ctx.info_name == 'archive':
            validate_past(value)
        return value
    except (TypeError, ValueError) as e:
        logger.debug("bad timestamp parameter", exc_info=True)
        raise click.BadParameter(unicode(e))

@main.command()
@click.option('--format', default='csv', type=click.Choice(['json', 'csv']),
              help='the output format')
@click.option('--begin', default=None, type=DateTime(),
              callback=to_timestamp,
              help='when to start archiving, as unix timestamp')
def archive(format, *args, **kwargs):
    dest = csv.Destination() if format == 'csv' else json.Destination()
    mgship.Archive(dest, *args, **kwargs).ship()

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>from __future__ import division
import os
import json
import tables
import shutil
import logging
from glob import glob
import numpy as np
import matplotlib
if __name__ == '__main__':
    matplotlib.use('Agg')
import matplotlib.pyplot as plt

import Ska.DBI

import mica.archive.asp_l1 as asp_l1_arch
import mica.archive.obspar as obspar_arch
from mica.archive import obsid_archive
from .core import Obi

# get rid of the black edges on the plot markers
plt.rcParams['lines.markeredgewidth'] = 0

mica_archive = os.environ.get('MICA_ARCHIVE') or '/data/aca/archive'
FILES = dict(
    data_root=os.path.join(mica_archive, 'vv'),
    temp_root=os.path.join(mica_archive, 'tempvv'),
    shelf_file=os.path.join(mica_archive, 'vv', 'vv_shelf.db'),
    h5_file=os.path.join(mica_archive, 'vv', 'vv.h5'),
    last_file=os.path.join(mica_archive, 'vv', 'last_id.txt'),
    asp1_proc_table=os.path.join(mica_archive, 'asp1', 'processing_asp_l1.db3'),
    bad_obsid_list=os.path.join(mica_archive, 'vv', 'bad_obsids.json'))

KNOWN_BAD_OBSIDS = json.loads(open(FILES['bad_obsid_list']).read())
#    save=True,
#    known_bad_obsids=[1283, 5542, 722, 721],

logger = logging.getLogger('vv')
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler())

def get_vv_dir(obsid, version="default"):
    """
    Get directory containing V&V products for a requested obsid/version,
    including plots and json.

    :param obsid: obsid
    :param version: 'last', 'default' or version number
    :returns: directory name for obsid/version
    """
    num_version = None
    if version == 'last' or version == 'default':
        asp_l1_proc = Ska.DBI.DBI(dbi="sqlite", server=FILES['asp1_proc_table'])
        if version == 'default':
            obs = asp_l1_proc.fetchall("""select * from aspect_1_proc
                                          where obsid = {} and isdefault = 1
                                       """.format(obsid))
            if not len(obs):
                raise ValueError("Version {} not found for obsid {}".format(
                    version, obsid))
            num_version = obs['revision'][0]
        if version == 'last':
            obs = asp_l1_proc.fetchall("""select * from aspect_1_proc
                                          where obsid = {}
                                       """.format(obsid))
            if not len(obs):
                raise ValueError("Version {} not found for obsid {}".format(
                    obsid, version))
            num_version = np.max(obs['revision'])
    else:
        num_version = version
    strobs = "%05d_v%02d" % (obsid, num_version)
    chunk_dir = strobs[0:2]
    chunk_dir_path = os.path.join(FILES['data_root'], chunk_dir)
    obs_dir = os.path.join(chunk_dir_path, strobs)
    return obs_dir

def get_vv_files(obsid, version="default"):
    """
    Get list of V&V files available for a requested obsid/version.

    :param obsid: obsid
    :param version: 'default', 'last' or version number
    :returns: list of files
    """
    vv_dir = get_vv_dir(obsid, version)
    return glob(os.path.join(vv_dir, "*"))

def get_vv(obsid, version="default"):
    """
    Retrieve V&V data for an obsid/version.
    This reads the saved JSON and returns the previously-
    calculated V&V data.

    :param obsid: obsid
    :param version: 'last', 'default', or version number
    :returns: dict of V&V data
    """
    vv_dir = get_vv_dir(obsid, version)
    json_file = glob(os.path.join(vv_dir, "*.json"))[0]
    return json.loads(open(json_file).read())

def get_rms_data():
    """
    Retrieve/return all data from RMS trending H5 archive

    :returns: numpy array of RMS data for each star/obsid/version
    """
    h5f = tables.openFile(FILES['h5_file'], 'r')
    tbl = h5f.getNode('/', 'vv')
    data = tbl[:]
    h5f.close()
    return data

def _file_vv(obi):
    """
    Save processed V&V data to per-obsid archive
    """
    obsid = int(obi.info()['obsid'])
    version = int(obi.info()['revision'])
    # set up directory for data
    strobs = "%05d_v%02d" % (obsid, version)
    chunk_dir = strobs[0:2]
    chunk_dir_path = os.path.join(FILES['data_root'], chunk_dir)
    obs_dir = os.path.join(chunk_dir_path, strobs)
    if not os.path.exists(obs_dir):
        logger.info("making directory %s" % obs_dir)
        os.makedirs(obs_dir)
    else:
        logger.info("obsid dir %s already exists" % obs_dir)
    for f in glob(os.path.join(obi.tempdir, "*")):
        os.chmod(f, 0775)
        shutil.copy(f, obs_dir)
        os.remove(f)
    logger.info("moved VV files to {}".format(obs_dir))
    os.removedirs(obi.tempdir)
    logger.info("removed directory {}".format(obi.tempdir))
    # make any desired link
    obs_ln = os.path.join(FILES['data_root'], chunk_dir, "%05d" % obsid)
    obs_ln_last = os.path.join(
        FILES['data_root'], chunk_dir, "%05d_last" % obsid)
    obsdirs = asp_l1_arch.get_obs_dirs(obsid)
    isdefault = 0
    if 'default' in obsdirs:
        if (os.path.realpath(obsdirs[version])
                == os.path.realpath(obsdirs['default'])):
            if os.path.islink(obs_ln):
                os.unlink(obs_ln)
            os.symlink(os.path.relpath(obs_dir, chunk_dir_path), obs_ln)
            isdefault = 1
    if 'last' in obsdirs:
        if ('default' in obsdirs
                and (os.path.realpath(obsdirs['last'])
                     != os.path.realpath(obsdirs['default']))
                or 'default' not in obsdirs):
            if (os.path.realpath(obsdirs[version])
                    == os.path.realpath(obsdirs['last'])):
                if os.path.islink(obs_ln_last):
                    os.unlink(obs_ln_last)
                os.symlink(os.path.relpath(obs_dir, chunk_dir_path),
                           obs_ln_last)
        if ('default' in obsdirs
            and (os.path.realpath(obsdirs['last'])
                 == os.path.realpath(obsdirs['default']))):
            if os.path.exists(obs_ln_last):
                os.unlink(obs_ln_last)
    obi.isdefault = isdefault

def update(obsids=[]):
    """
    For a list of obsids or for all new obsids, run V&V processing
    and save V&V info to archive.

    :param obsids: optional list of obsids
    """
    asp_l1_proc = Ska.DBI.DBI(dbi="sqlite", server=FILES['asp1_proc_table'])
    # if an obsid is requested, just do that
    # there's a little bit of duplication in this block
    if len(obsids):
        for obsid in obsids:
            todo = asp_l1_proc.fetchall(
                "SELECT * FROM aspect_1_proc where obsid = {}".format(
                    obsid))
            for obs in todo:
                if obs['obsid'] in KNOWN_BAD_OBSIDS:
                    logger.info("Skipping known bad obsid {}".format(obs['obsid']))
                    continue
                logger.info("running VV for obsid {} run on {}".format(
                    obs['obsid'], obs['ap_date']))
                process(obs['obsid'], version=obs['revision'])
                asp_l1_proc.execute("""UPDATE aspect_1_proc set vv_complete = 1
                                       where obsid = {} and revision = {}
                                    """.format(obs['obsid'], obs['revision']))
    else:
        # if no obsid specified, try to retrieve all data without vv
        todo = asp_l1_proc.fetchall(
            """SELECT * FROM aspect_1_proc
               where vv_complete != 1 order by aspect_1_id""")
        for obs in todo:
            if obs['obsid'] in KNOWN_BAD_OBSIDS:
                logger.info("Skipping known bad obsid {}".format(obs['obsid']))
                continue
            logger.info("running VV for obsid {} run on {}".format(
                obs['obsid'], obs['ap_date']))
            process(obs['obsid'], version=obs['revision'])
            asp_l1_proc.execute("""UPDATE aspect_1_proc set vv_complete = 1
                                   where obsid = {} and revision = {}
                                """.format(obs['obsid'], obs['revision']))

def get_arch_vv(obsid, version='last'):
    """
    Given obsid and version, find archived ASP1 and obspar products and
    run V&V.  Effort is made to find the obspar that was actually used during
    creation of the ASP1 products.

    :param obsid: obsid
    :param version: 'last', 'default', or revision number of ASP1 products
    :returns: mica.vv.Obi V&V object
    """
    asp_l1_dirs = asp_l1_arch.get_obs_dirs(obsid)
    if version not in asp_l1_dirs:
        logger.error(
            "Requested version {} not in asp_l1 archive".format(version))
        return None
    l1_dir = asp_l1_dirs[version]
    # find the obspar that matches the requested aspect_1 products
    # this is in the aspect processing table
    asp_l1_proc = Ska.DBI.DBI(dbi="sqlite", server=FILES['asp1_proc_table'])
    asp_obs = asp_l1_proc.fetchall(
        "SELECT * FROM aspect_1_proc where obsid = {}".format(
            obsid))
    asp_proc = None
    if version == 'last':
        asp_proc = asp_obs[asp_obs['aspect_1_id']
                           == np.max(asp_obs['aspect_1_id'])][0]
    if version == 'default':
        asp_proc = asp_obs[asp_obs['isdefault'] == 1][0]
    if asp_proc is None:
        asp_proc = asp_obs[asp_obs['revision'] == version][0]
    obspar_dirs = obspar_arch.get_obs_dirs(obsid)
    if asp_proc['obspar_version'] not in obspar_dirs:
        # try to update the obspar archive with the missing version
        config = obspar_arch.CONFIG.copy()
        config.update(dict(obsid=obsid, version=asp_proc['obspar_version']))
        oa = obsid_archive.ObsArchive(config)
        oa.logger.setLevel(logging.INFO)
        oa.logger.addHandler(logging.StreamHandler())
        oa.update()
        obspar_dirs = obspar_arch.get_obs_dirs(obsid)
    obspar_file = glob(os.path.join(obspar_dirs[asp_proc['obspar_version']],
                                    'axaf*par*'))[0]
    return Obi(obspar_file, l1_dir, temproot=FILES['temp_root'])

def process(obsid, version='last'):
    """
    For requested obsid/version, run V&V, make plots,
    save plots and JSON, save info to shelve file, and
    update RMS HDF5.

    :param obsid: obsid
    :param version: 'last', 'default' or revision number of ASP1 products
    :returns: mica.vv.Obi V&V object
    """
    obi = get_arch_vv(obsid, version)
    obi.save_plots_and_resid()
    if obi is None:
        return None
    shelf_file = os.path.join(FILES['data_root'],
                              FILES['shelf_file'])
    obi.shelve_info(shelf_file)
    _file_vv(obi)
    h5 = tables.openFile(FILES['h5_file'], 'a')
    tbl = h5.getNode('/', 'vv')
    obi.set_tbl(tbl)
    obi.slots_to_table()
    tbl.flush()
    h5.flush()
    h5.close()
    return obi

<commit_msg>PEP8 fixes for vv<commit_after>from __future__ import division
import os
import json
import tables
import shutil
import logging
from glob import glob
import numpy as np
import matplotlib
if __name__ == '__main__':
    matplotlib.use('Agg')
import matplotlib.pyplot as plt

import Ska.DBI

import mica.archive.asp_l1 as asp_l1_arch
import mica.archive.obspar as obspar_arch
from mica.archive import obsid_archive
from .core import Obi

# get rid of the black edges on the plot markers
plt.rcParams['lines.markeredgewidth'] = 0

mica_archive = os.environ.get('MICA_ARCHIVE') or '/data/aca/archive'
FILES = dict(
    data_root=os.path.join(mica_archive, 'vv'),
    temp_root=os.path.join(mica_archive, 'tempvv'),
    shelf_file=os.path.join(mica_archive, 'vv', 'vv_shelf.db'),
    h5_file=os.path.join(mica_archive, 'vv', 'vv.h5'),
    last_file=os.path.join(mica_archive, 'vv', 'last_id.txt'),
    asp1_proc_table=os.path.join(mica_archive, 'asp1', 'processing_asp_l1.db3'),
    bad_obsid_list=os.path.join(mica_archive, 'vv', 'bad_obsids.json'))

KNOWN_BAD_OBSIDS = json.loads(open(FILES['bad_obsid_list']).read())
#    save=True,
#    known_bad_obsids=[1283, 5542, 722, 721],

logger = logging.getLogger('vv')
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler())

def get_vv_dir(obsid, version="default"):
    """
    Get directory containing V&V products for a requested obsid/version,
    including plots and json.

    :param obsid: obsid
    :param version: 'last', 'default' or version number
    :returns: directory name for obsid/version
    """
    num_version = None
    if version == 'last' or version == 'default':
        asp_l1_proc = Ska.DBI.DBI(dbi="sqlite", server=FILES['asp1_proc_table'])
        if version == 'default':
            obs = asp_l1_proc.fetchall("""select * from aspect_1_proc
                                          where obsid = {} and isdefault = 1
                                       """.format(obsid))
            if not len(obs):
                raise ValueError("Version {} not found for obsid {}".format(
                    version, obsid))
            num_version = obs['revision'][0]
        if version == 'last':
            obs = asp_l1_proc.fetchall("""select * from aspect_1_proc
                                          where obsid = {}
                                       """.format(obsid))
            if not len(obs):
                raise ValueError("Version {} not found for obsid {}".format(
                    obsid, version))
            num_version = np.max(obs['revision'])
    else:
        num_version = version
    strobs = "%05d_v%02d" % (obsid, num_version)
    chunk_dir = strobs[0:2]
    chunk_dir_path = os.path.join(FILES['data_root'], chunk_dir)
    obs_dir = os.path.join(chunk_dir_path, strobs)
    return obs_dir

def get_vv_files(obsid, version="default"):
    """
    Get list of V&V files available for a requested obsid/version.

    :param obsid: obsid
    :param version: 'default', 'last' or version number
    :returns: list of files
    """
    vv_dir = get_vv_dir(obsid, version)
    return glob(os.path.join(vv_dir, "*"))

def get_vv(obsid, version="default"):
    """
    Retrieve V&V data for an obsid/version.
    This reads the saved JSON and returns the previously-
    calculated V&V data.

    :param obsid: obsid
    :param version: 'last', 'default', or version number
    :returns: dict of V&V data
    """
    vv_dir = get_vv_dir(obsid, version)
    json_file = glob(os.path.join(vv_dir, "*.json"))[0]
    return json.loads(open(json_file).read())

def get_rms_data():
    """
    Retrieve/return all data from RMS trending H5 archive

    :returns: numpy array of RMS data for each star/obsid/version
    """
    h5f = tables.openFile(FILES['h5_file'], 'r')
    tbl = h5f.getNode('/', 'vv')
    data = tbl[:]
    h5f.close()
    return data

def _file_vv(obi):
    """
    Save processed V&V data to per-obsid archive
    """
    obsid = int(obi.info()['obsid'])
    version = int(obi.info()['revision'])
    # set up directory for data
    strobs = "%05d_v%02d" % (obsid, version)
    chunk_dir = strobs[0:2]
    chunk_dir_path = os.path.join(FILES['data_root'], chunk_dir)
    obs_dir = os.path.join(chunk_dir_path, strobs)
    if not os.path.exists(obs_dir):
        logger.info("making directory %s" % obs_dir)
        os.makedirs(obs_dir)
    else:
        logger.info("obsid dir %s already exists" % obs_dir)
    for f in glob(os.path.join(obi.tempdir, "*")):
        os.chmod(f, 0775)
        shutil.copy(f, obs_dir)
        os.remove(f)
    logger.info("moved VV files to {}".format(obs_dir))
    os.removedirs(obi.tempdir)
    logger.info("removed directory {}".format(obi.tempdir))
    # make any desired link
    obs_ln = os.path.join(FILES['data_root'], chunk_dir, "%05d" % obsid)
    obs_ln_last = os.path.join(
        FILES['data_root'], chunk_dir, "%05d_last" % obsid)
    obsdirs = asp_l1_arch.get_obs_dirs(obsid)
    isdefault = 0
    if 'default' in obsdirs:
        if (os.path.realpath(obsdirs[version])
                == os.path.realpath(obsdirs['default'])):
            if os.path.islink(obs_ln):
                os.unlink(obs_ln)
            os.symlink(os.path.relpath(obs_dir, chunk_dir_path), obs_ln)
            isdefault = 1
    if 'last' in obsdirs:
        if ('default' in obsdirs
                and (os.path.realpath(obsdirs['last'])
                     != os.path.realpath(obsdirs['default']))
                or 'default' not in obsdirs):
            if (os.path.realpath(obsdirs[version])
                    == os.path.realpath(obsdirs['last'])):
                if os.path.islink(obs_ln_last):
                    os.unlink(obs_ln_last)
                os.symlink(os.path.relpath(obs_dir, chunk_dir_path),
                           obs_ln_last)
        if ('default' in obsdirs
            and (os.path.realpath(obsdirs['last'])
                 == os.path.realpath(obsdirs['default']))):
            if os.path.exists(obs_ln_last):
                os.unlink(obs_ln_last)
    obi.isdefault = isdefault

def update(obsids=[]):
    """
    For a list of obsids or for all new obsids, run V&V processing
    and save V&V info to archive.

    :param obsids: optional list of obsids
    """
    asp_l1_proc = Ska.DBI.DBI(dbi="sqlite", server=FILES['asp1_proc_table'])
    # if an obsid is requested, just do that
    # there's a little bit of duplication in this block
    if len(obsids):
        for obsid in obsids:
            todo = asp_l1_proc.fetchall(
                "SELECT * FROM aspect_1_proc where obsid = {}".format(
                    obsid))
            for obs in todo:
                if obs['obsid'] in KNOWN_BAD_OBSIDS:
                    logger.info("Skipping known bad obsid {}".format(obs['obsid']))
                    continue
                logger.info("running VV for obsid {} run on {}".format(
                    obs['obsid'], obs['ap_date']))
                process(obs['obsid'], version=obs['revision'])
                asp_l1_proc.execute("""UPDATE aspect_1_proc set vv_complete = 1
                                       where obsid = {} and revision = {}
                                    """.format(obs['obsid'], obs['revision']))
    else:
        # if no obsid specified, try to retrieve all data without vv
        todo = asp_l1_proc.fetchall(
            """SELECT * FROM aspect_1_proc
               where vv_complete != 1 order by aspect_1_id""")
        for obs in todo:
            if obs['obsid'] in KNOWN_BAD_OBSIDS:
                logger.info("Skipping known bad obsid {}".format(obs['obsid']))
                continue
            logger.info("running VV for obsid {} run on {}".format(
                obs['obsid'], obs['ap_date']))
            process(obs['obsid'], version=obs['revision'])
            asp_l1_proc.execute("""UPDATE aspect_1_proc set vv_complete = 1
                                   where obsid = {} and revision = {}
                                """.format(obs['obsid'], obs['revision']))

def get_arch_vv(obsid, version='last'):
    """
    Given obsid and version, find archived ASP1 and obspar products and
    run V&V.  Effort is made to find the obspar that was actually used during
    creation of the ASP1 products.

    :param obsid: obsid
    :param version: 'last', 'default', or revision number of ASP1 products
    :returns: mica.vv.Obi V&V object
    """
    asp_l1_dirs = asp_l1_arch.get_obs_dirs(obsid)
    if version not in asp_l1_dirs:
        logger.error(
            "Requested version {} not in asp_l1 archive".format(version))
        return None
    l1_dir = asp_l1_dirs[version]
    # find the obspar that matches the requested aspect_1 products
    # this is in the aspect processing table
    asp_l1_proc = Ska.DBI.DBI(dbi="sqlite", server=FILES['asp1_proc_table'])
    asp_obs = asp_l1_proc.fetchall(
        "SELECT * FROM aspect_1_proc where obsid = {}".format(
            obsid))
    asp_proc = None
    if version == 'last':
        asp_proc = asp_obs[asp_obs['aspect_1_id']
                           == np.max(asp_obs['aspect_1_id'])][0]
    if version == 'default':
        asp_proc = asp_obs[asp_obs['isdefault'] == 1][0]
    if asp_proc is None:
        asp_proc = asp_obs[asp_obs['revision'] == version][0]
    obspar_dirs = obspar_arch.get_obs_dirs(obsid)
    if asp_proc['obspar_version'] not in obspar_dirs:
        # try to update the obspar archive with the missing version
        config = obspar_arch.CONFIG.copy()
        config.update(dict(obsid=obsid, version=asp_proc['obspar_version']))
        oa = obsid_archive.ObsArchive(config)
        oa.logger.setLevel(logging.INFO)
        oa.logger.addHandler(logging.StreamHandler())
        oa.update()
        obspar_dirs = obspar_arch.get_obs_dirs(obsid)
    obspar_file = glob(os.path.join(obspar_dirs[asp_proc['obspar_version']],
                                    'axaf*par*'))[0]
    return Obi(obspar_file, l1_dir, temproot=FILES['temp_root'])

def process(obsid, version='last'):
    """
    For requested obsid/version, run V&V, make plots,
    save plots and JSON, save info to shelve file, and
    update RMS HDF5.

    :param obsid: obsid
    :param version: 'last', 'default' or revision number of ASP1 products
    :returns: mica.vv.Obi V&V object
    """
    obi = get_arch_vv(obsid, version)
    obi.save_plots_and_resid()
    if obi is None:
        return None
    shelf_file = os.path.join(FILES['data_root'],
                              FILES['shelf_file'])
    obi.shelve_info(shelf_file)
    _file_vv(obi)
    h5 = tables.openFile(FILES['h5_file'], 'a')
    tbl = h5.getNode('/', 'vv')
    obi.set_tbl(tbl)
    obi.slots_to_table()
    tbl.flush()
    h5.flush()
    h5.close()
    return obi
<|endoftext|>


<commit_before>import sys
import os

from . import options
from . import builtin_functions
from . import loading
from . import utilities
from . import error
from . import file_selection

def main():
    try:
        options.initialize_environment_variables()

        processed_options = options.process_options()
        filename = processed_options.script
        if filename != '-':
            filename = file_selection.try_select_path(
                os.curdir,
                os.curdir,
                filename,
            )

        loading.load_file(filename, {
            **builtin_functions.BUILTIN_FUNCTIONS,
            **options.make_args_function(processed_options),
        }, processed_options.target, utilities.get_base_path(filename))
    except Exception as exception:
        sys.exit('error: ' + str(exception))
    except KeyboardInterrupt:
        # output a line break after the ^C symbol in a terminal
        print('')

        sys.exit(1)
<commit_msg>Issue #127: create a specific file cache<commit_after>import sys
import os

from . import options
from . import builtin_functions
from . import loading
from . import utilities
from . import error
from . import file_selection
from . import file_cache

def main():
    try:
        options.initialize_environment_variables()

        processed_options = options.process_options()
        filename = processed_options.script
        if filename != '-':
            filename = file_selection.try_select_path(
                os.curdir,
                os.curdir,
                filename,
            )

        specific_file_cache = file_cache.FileCache()
        loading.load_file(
            filename,
            {
                **builtin_functions.BUILTIN_FUNCTIONS,
                **options.make_args_function(processed_options),
            },
            processed_options.target,
            utilities.get_base_path(filename),
            specific_file_cache,
        )
    except Exception as exception:
        sys.exit('error: ' + str(exception))
    except KeyboardInterrupt:
        # output a line break after the ^C symbol in a terminal
        print('')

        sys.exit(1)
<|endoftext|>


<commit_before>"""
This module implements a transaction manager that can be used to define
transaction handling in a request or view function. It is used by transaction
control middleware and decorators.

The transaction manager can be in managed or in auto state. Auto state means the
system is using a commit-on-save strategy (actually it's more like
commit-on-change). As soon as the .save() or .delete() (or related) methods are
called, a commit is made.

Managed transactions don't do those commits, but will need some kind of manual
or implicit commits or rollbacks.
"""

try:
    import thread
except ImportError:
    import dummy_thread as thread
try:
    from functools import wraps
except ImportError:
    from django.utils.functional import wraps  # Python 2.3, 2.4 fallback.
from django.db import connections, DEFAULT_DB_ALIAS
from django.conf import settings

class TransactionManagementError(Exception):
    """
    This exception is thrown when something bad happens with transaction
    management.
    """
    pass

# The states are dictionaries of dictionaries of lists. The key to the outer
# dict is the current thread, and the key to the inner dictionary is the
# connection alias and the list is handled as a stack of values.
state = {}
savepoint_state = {}

# The dirty flag is set by *_unless_managed functions to denote that the
# code under transaction management has changed things to require a
# database commit.
# This is a dictionary mapping thread to a dictionary mapping connection
# alias to a boolean.
dirty = {}

def enter_transaction_management(managed=True, using=None):
    """
    Enters transaction management for a running thread. It must be balanced with
    the appropriate leave_transaction_management call, since the actual state is
    managed as a stack.

    The state and dirty flag are carried over from the surrounding block or
    from the settings, if there is no surrounding block (dirty is always false
    when no current block is running).
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in state and state[thread_ident].get(using):
        state[thread_ident][using].append(state[thread_ident][using][-1])
    else:
        state.setdefault(thread_ident, {})
        state[thread_ident][using] = [settings.TRANSACTIONS_MANAGED]
    if thread_ident not in dirty or using not in dirty[thread_ident]:
        dirty.setdefault(thread_ident, {})
        dirty[thread_ident][using] = False
    connection._enter_transaction_management(managed)

def leave_transaction_management(using=None):
    """
    Leaves transaction management for a running thread. A dirty flag is carried
    over to the surrounding block, as a commit will commit all changes, even
    those from outside. (Commits are on connection level.)
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    connection._leave_transaction_management(is_managed(using=using))
    thread_ident = thread.get_ident()
    if thread_ident in state and state[thread_ident].get(using):
        del state[thread_ident][using][-1]
    else:
        raise TransactionManagementError("This code isn't under transaction management")
    if dirty.get(thread_ident, {}).get(using, False):
        rollback(using=using)
        raise TransactionManagementError("Transaction managed block ended with pending COMMIT/ROLLBACK")
    dirty[thread_ident][using] = False

def is_dirty(using=None):
    """
    Returns True if the current transaction requires a commit for changes to
    happen.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    return dirty.get(thread.get_ident(), {}).get(using, False)

def set_dirty(using=None):
    """
    Sets a dirty flag for the current thread and code streak. This can be used
    to decide in a managed block of code to decide whether there are open
    changes waiting for commit.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in dirty and using in dirty[thread_ident]:
        dirty[thread_ident][using] = True
    else:
        raise TransactionManagementError("This code isn't under transaction management")

def set_clean(using=None):
    """
    Resets a dirty flag for the current thread and code streak. This can be used
    to decide in a managed block of code to decide whether a commit or rollback
    should happen.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in dirty and using in dirty[thread_ident]:
        dirty[thread_ident][using] = False
    else:
        raise TransactionManagementError("This code isn't under transaction management")
    clean_savepoints(using=using)

def clean_savepoints(using=None):
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        del savepoint_state[thread_ident][using]

def is_managed(using=None):
    """
    Checks whether the transaction manager is in manual or in auto state.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in state and using in state[thread_ident]:
        if state[thread_ident][using]:
            return state[thread_ident][using][-1]
    return settings.TRANSACTIONS_MANAGED

def managed(flag=True, using=None):
    """
    Puts the transaction manager into a manual state: managed transactions have
    to be committed explicitly by the user. If you switch off transaction
    management and there is a pending commit/rollback, the data will be
    commited.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    top = state.get(thread_ident, {}).get(using, None)
    if top:
        top[-1] = flag
        if not flag and is_dirty(using=using):
            connection._commit()
            set_clean(using=using)
    else:
        raise TransactionManagementError("This code isn't under transaction management")

def commit_unless_managed(using=None):
    """
    Commits changes if the system is not in managed transaction mode.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    if not is_managed(using=using):
        connection._commit()
        clean_savepoints(using=using)
    else:
        set_dirty(using=using)

def rollback_unless_managed(using=None):
    """
    Rolls back changes if the system is not in managed transaction mode.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    if not is_managed(using=using):
        connection._rollback()
    else:
        set_dirty(using=using)

def commit(using=None):
    """
    Does the commit itself and resets the dirty flag.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    connection._commit()
    set_clean(using=using)

def rollback(using=None):
    """
    This function does the rollback itself and resets the dirty flag.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    connection._rollback()
    set_clean(using=using)

def savepoint(using=None):
    """
    Creates a savepoint (if supported and required by the backend) inside the
    current transaction. Returns an identifier for the savepoint that will be
    used for the subsequent rollback or commit.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        savepoint_state[thread_ident][using].append(None)
    else:
        savepoint_state.setdefault(thread_ident, {})
        savepoint_state[thread_ident][using] = [None]
    tid = str(thread_ident).replace('-', '')
    sid = "s%s_x%d" % (tid, len(savepoint_state[thread_ident][using]))
    connection._savepoint(sid)
    return sid

def savepoint_rollback(sid, using=None):
    """
    Rolls back the most recent savepoint (if one exists). Does nothing if
    savepoints are not supported.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        connection._savepoint_rollback(sid)

def savepoint_commit(sid, using=None):
    """
    Commits the most recent savepoint (if one exists). Does nothing if
    savepoints are not supported.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        connection._savepoint_commit(sid)

##############
# DECORATORS #
##############

def autocommit(func_or_using=None):
    """
    Decorator that activates commit on save. This is Django's default behavior;
    this decorator is useful if you globally activated transaction management in
    your settings file and want the default behavior in some view functions.
    """
    def inner_autocommit(func, using=None):
        def _autocommit(*args, **kw):
            try:
                enter_transaction_management(managed=False, using=using)
                managed(False, using=using)
                return func(*args, **kw)
            finally:
                leave_transaction_management(using=using)
        return wraps(func)(_autocommit)
    if func_or_using is None:
        func_or_using = DEFAULT_DB_ALIAS
    if callable(func_or_using):
        return inner_autocommit(func_or_using, DEFAULT_DB_ALIAS)
    return lambda func: inner_autocommit(func,  func_or_using)

def commit_on_success(func_or_using=None):
    """
    This decorator activates commit on response. This way, if the view function
    runs successfully, a commit is made; if the viewfunc produces an exception,
    a rollback is made. This is one of the most common ways to do transaction
    control in web apps.
    """
    def inner_commit_on_success(func, using=None):
        def _commit_on_success(*args, **kw):
            try:
                enter_transaction_management(using=using)
                managed(True, using=using)
                try:
                    res = func(*args, **kw)
                except:
                    # All exceptions must be handled here (even string ones).
                    if is_dirty(using=using):
                        rollback(using=using)
                    raise
                else:
                    if is_dirty(using=using):
                        commit(using=using)
                return res
            finally:
                leave_transaction_management(using=using)
        return wraps(func)(_commit_on_success)
    if func_or_using is None:
        func_or_using = DEFAULT_DB_ALIAS
    if callable(func_or_using):
        return inner_commit_on_success(func_or_using, DEFAULT_DB_ALIAS)
    return lambda func: inner_commit_on_success(func, func_or_using)

def commit_manually(func_or_using=None):
    """
    Decorator that activates manual transaction control. It just disables
    automatic transaction control and doesn't do any commit/rollback of its
    own -- it's up to the user to call the commit and rollback functions
    themselves.
    """
    def inner_commit_manually(func, using=None):
        def _commit_manually(*args, **kw):
            try:
                enter_transaction_management(using=using)
                managed(True, using=using)
                return func(*args, **kw)
            finally:
                leave_transaction_management(using=using)

        return wraps(func)(_commit_manually)
    if func_or_using is None:
        func_or_using = DEFALUT_DB_ALIAS
    if callable(func_or_using):
        return inner_commit_manually(func_or_using, DEFAULT_DB_ALIAS)
    return lambda func: inner_commit_manually(func, func_or_using)
<commit_msg>Fixed typo introduced in r11952.<commit_after>"""
This module implements a transaction manager that can be used to define
transaction handling in a request or view function. It is used by transaction
control middleware and decorators.

The transaction manager can be in managed or in auto state. Auto state means the
system is using a commit-on-save strategy (actually it's more like
commit-on-change). As soon as the .save() or .delete() (or related) methods are
called, a commit is made.

Managed transactions don't do those commits, but will need some kind of manual
or implicit commits or rollbacks.
"""

try:
    import thread
except ImportError:
    import dummy_thread as thread
try:
    from functools import wraps
except ImportError:
    from django.utils.functional import wraps  # Python 2.3, 2.4 fallback.
from django.db import connections, DEFAULT_DB_ALIAS
from django.conf import settings

class TransactionManagementError(Exception):
    """
    This exception is thrown when something bad happens with transaction
    management.
    """
    pass

# The states are dictionaries of dictionaries of lists. The key to the outer
# dict is the current thread, and the key to the inner dictionary is the
# connection alias and the list is handled as a stack of values.
state = {}
savepoint_state = {}

# The dirty flag is set by *_unless_managed functions to denote that the
# code under transaction management has changed things to require a
# database commit.
# This is a dictionary mapping thread to a dictionary mapping connection
# alias to a boolean.
dirty = {}

def enter_transaction_management(managed=True, using=None):
    """
    Enters transaction management for a running thread. It must be balanced with
    the appropriate leave_transaction_management call, since the actual state is
    managed as a stack.

    The state and dirty flag are carried over from the surrounding block or
    from the settings, if there is no surrounding block (dirty is always false
    when no current block is running).
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in state and state[thread_ident].get(using):
        state[thread_ident][using].append(state[thread_ident][using][-1])
    else:
        state.setdefault(thread_ident, {})
        state[thread_ident][using] = [settings.TRANSACTIONS_MANAGED]
    if thread_ident not in dirty or using not in dirty[thread_ident]:
        dirty.setdefault(thread_ident, {})
        dirty[thread_ident][using] = False
    connection._enter_transaction_management(managed)

def leave_transaction_management(using=None):
    """
    Leaves transaction management for a running thread. A dirty flag is carried
    over to the surrounding block, as a commit will commit all changes, even
    those from outside. (Commits are on connection level.)
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    connection._leave_transaction_management(is_managed(using=using))
    thread_ident = thread.get_ident()
    if thread_ident in state and state[thread_ident].get(using):
        del state[thread_ident][using][-1]
    else:
        raise TransactionManagementError("This code isn't under transaction management")
    if dirty.get(thread_ident, {}).get(using, False):
        rollback(using=using)
        raise TransactionManagementError("Transaction managed block ended with pending COMMIT/ROLLBACK")
    dirty[thread_ident][using] = False

def is_dirty(using=None):
    """
    Returns True if the current transaction requires a commit for changes to
    happen.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    return dirty.get(thread.get_ident(), {}).get(using, False)

def set_dirty(using=None):
    """
    Sets a dirty flag for the current thread and code streak. This can be used
    to decide in a managed block of code to decide whether there are open
    changes waiting for commit.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in dirty and using in dirty[thread_ident]:
        dirty[thread_ident][using] = True
    else:
        raise TransactionManagementError("This code isn't under transaction management")

def set_clean(using=None):
    """
    Resets a dirty flag for the current thread and code streak. This can be used
    to decide in a managed block of code to decide whether a commit or rollback
    should happen.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in dirty and using in dirty[thread_ident]:
        dirty[thread_ident][using] = False
    else:
        raise TransactionManagementError("This code isn't under transaction management")
    clean_savepoints(using=using)

def clean_savepoints(using=None):
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        del savepoint_state[thread_ident][using]

def is_managed(using=None):
    """
    Checks whether the transaction manager is in manual or in auto state.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    thread_ident = thread.get_ident()
    if thread_ident in state and using in state[thread_ident]:
        if state[thread_ident][using]:
            return state[thread_ident][using][-1]
    return settings.TRANSACTIONS_MANAGED

def managed(flag=True, using=None):
    """
    Puts the transaction manager into a manual state: managed transactions have
    to be committed explicitly by the user. If you switch off transaction
    management and there is a pending commit/rollback, the data will be
    commited.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    top = state.get(thread_ident, {}).get(using, None)
    if top:
        top[-1] = flag
        if not flag and is_dirty(using=using):
            connection._commit()
            set_clean(using=using)
    else:
        raise TransactionManagementError("This code isn't under transaction management")

def commit_unless_managed(using=None):
    """
    Commits changes if the system is not in managed transaction mode.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    if not is_managed(using=using):
        connection._commit()
        clean_savepoints(using=using)
    else:
        set_dirty(using=using)

def rollback_unless_managed(using=None):
    """
    Rolls back changes if the system is not in managed transaction mode.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    if not is_managed(using=using):
        connection._rollback()
    else:
        set_dirty(using=using)

def commit(using=None):
    """
    Does the commit itself and resets the dirty flag.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    connection._commit()
    set_clean(using=using)

def rollback(using=None):
    """
    This function does the rollback itself and resets the dirty flag.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    connection._rollback()
    set_clean(using=using)

def savepoint(using=None):
    """
    Creates a savepoint (if supported and required by the backend) inside the
    current transaction. Returns an identifier for the savepoint that will be
    used for the subsequent rollback or commit.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        savepoint_state[thread_ident][using].append(None)
    else:
        savepoint_state.setdefault(thread_ident, {})
        savepoint_state[thread_ident][using] = [None]
    tid = str(thread_ident).replace('-', '')
    sid = "s%s_x%d" % (tid, len(savepoint_state[thread_ident][using]))
    connection._savepoint(sid)
    return sid

def savepoint_rollback(sid, using=None):
    """
    Rolls back the most recent savepoint (if one exists). Does nothing if
    savepoints are not supported.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        connection._savepoint_rollback(sid)

def savepoint_commit(sid, using=None):
    """
    Commits the most recent savepoint (if one exists). Does nothing if
    savepoints are not supported.
    """
    if using is None:
        using = DEFAULT_DB_ALIAS
    connection = connections[using]
    thread_ident = thread.get_ident()
    if thread_ident in savepoint_state and using in savepoint_state[thread_ident]:
        connection._savepoint_commit(sid)

##############
# DECORATORS #
##############

def autocommit(func_or_using=None):
    """
    Decorator that activates commit on save. This is Django's default behavior;
    this decorator is useful if you globally activated transaction management in
    your settings file and want the default behavior in some view functions.
    """
    def inner_autocommit(func, using=None):
        def _autocommit(*args, **kw):
            try:
                enter_transaction_management(managed=False, using=using)
                managed(False, using=using)
                return func(*args, **kw)
            finally:
                leave_transaction_management(using=using)
        return wraps(func)(_autocommit)
    if func_or_using is None:
        func_or_using = DEFAULT_DB_ALIAS
    if callable(func_or_using):
        return inner_autocommit(func_or_using, DEFAULT_DB_ALIAS)
    return lambda func: inner_autocommit(func,  func_or_using)

def commit_on_success(func_or_using=None):
    """
    This decorator activates commit on response. This way, if the view function
    runs successfully, a commit is made; if the viewfunc produces an exception,
    a rollback is made. This is one of the most common ways to do transaction
    control in web apps.
    """
    def inner_commit_on_success(func, using=None):
        def _commit_on_success(*args, **kw):
            try:
                enter_transaction_management(using=using)
                managed(True, using=using)
                try:
                    res = func(*args, **kw)
                except:
                    # All exceptions must be handled here (even string ones).
                    if is_dirty(using=using):
                        rollback(using=using)
                    raise
                else:
                    if is_dirty(using=using):
                        commit(using=using)
                return res
            finally:
                leave_transaction_management(using=using)
        return wraps(func)(_commit_on_success)
    if func_or_using is None:
        func_or_using = DEFAULT_DB_ALIAS
    if callable(func_or_using):
        return inner_commit_on_success(func_or_using, DEFAULT_DB_ALIAS)
    return lambda func: inner_commit_on_success(func, func_or_using)

def commit_manually(func_or_using=None):
    """
    Decorator that activates manual transaction control. It just disables
    automatic transaction control and doesn't do any commit/rollback of its
    own -- it's up to the user to call the commit and rollback functions
    themselves.
    """
    def inner_commit_manually(func, using=None):
        def _commit_manually(*args, **kw):
            try:
                enter_transaction_management(using=using)
                managed(True, using=using)
                return func(*args, **kw)
            finally:
                leave_transaction_management(using=using)

        return wraps(func)(_commit_manually)
    if func_or_using is None:
        func_or_using = DEFAULT_DB_ALIAS
    if callable(func_or_using):
        return inner_commit_manually(func_or_using, DEFAULT_DB_ALIAS)
    return lambda func: inner_commit_manually(func, func_or_using)
<|endoftext|>